Add documentation to annot, chc, and pretty modules

src/annot.rs

@@ -1,3 +1,13 @@
+//! A parser for refinement type and formula annotations.
+//!
+//! This module provides a parser for `#[thrust::...]` attributes in Thrust. The parser is
+//! parameterized by the [`Resolver`] trait, which abstracts over the resolution of variable
+//! names. This allows the parser to be used in different contexts with different sets of
+//! available variables.
+//!
+//! The main entry point is [`AnnotParser`], which parses a [`TokenStream`] into a
+//! [`rty::RefinedType`] or a [`chc::Formula`].
+
 use rustc_ast::token::{BinOpToken, Delimiter, LitKind, Token, TokenKind};
 use rustc_ast::tokenstream::{RefTokenTreeCursor, Spacing, TokenStream, TokenTree};
 use rustc_index::IndexVec;
@@ -7,6 +17,9 @@ use crate::chc;
 use crate::pretty::PrettyDisplayExt as _;
 use crate::rty;
 
+/// A formula in an annotation.
+///
+/// This can be either a logical formula or the special value `auto` which tells Thrust to infer it.
 #[derive(Debug, Clone)]
 pub enum AnnotFormula<T> {
     Auto,
@@ -19,6 +32,7 @@ impl<T> AnnotFormula<T> {
     }
 }
 
+/// A trait for resolving variables in annotations to their logical representation and their sorts.
 pub trait Resolver {
     type Output;
     fn resolve(&self, ident: Ident) -> Option<(Self::Output, chc::Sort)>;
@@ -56,6 +70,7 @@ pub trait ResolverExt: Resolver {
 
 impl<T> ResolverExt for T where T: Resolver {}
 
+/// An error that can occur when parsing an attribute.
 #[derive(Debug, Clone, thiserror::Error)]
 pub enum ParseAttrError {
     #[error("unexpected end of input (expected {expected:?})")]
@@ -96,6 +111,7 @@ impl ParseAttrError {
 
 type Result<T> = std::result::Result<T, ParseAttrError>;
 
+/// A parser for refinement type annotations and formula annotations.
 struct Parser<'a, T> {
     resolver: T,
     cursor: RefTokenTreeCursor<'a>,
@@ -610,6 +626,7 @@ where
     }
 }
 
+/// A [`Resolver`] implementation for resolving specific variable as [`rty::RefinedTypeVar::Value`].
 struct RefinementResolver<'a, T> {
     resolver: Box<dyn Resolver<Output = T> + 'a>,
     self_: Option<(Ident, chc::Sort)>,
@@ -643,6 +660,7 @@ impl<'a, T> RefinementResolver<'a, T> {
     }
 }
 
+/// A [`Resolver`] that maps the output of another [`Resolver`].
 pub struct MappedResolver<'a, T, F> {
     resolver: Box<dyn Resolver<Output = T> + 'a>,
     map: F,
@@ -669,6 +687,9 @@ impl<'a, T, F> MappedResolver<'a, T, F> {
     }
 }
 
+/// A [`Resolver`] that stacks multiple [`Resolver`]s.
+///
+/// This resolver tries to resolve an identifier by querying a list of resolvers in order.
 pub struct StackedResolver<'a, T> {
     resolvers: Vec<Box<dyn Resolver<Output = T> + 'a>>,
 }
@@ -698,6 +719,7 @@ impl<'a, T> StackedResolver<'a, T> {
     }
 }
 
+/// A parser for annotations.
 #[derive(Debug, Clone)]
 pub struct AnnotParser<T> {
     resolver: T,

src/chc.rs

@@ -1,4 +1,5 @@
-/// Multi-sorted CHC system with tuples.
+//! A multi-sorted CHC system with tuples.
+
 use pretty::{termcolor, Pretty};
 use rustc_index::IndexVec;
 
@@ -17,6 +18,7 @@ pub use debug::DebugInfo;
 pub use solver::{CheckSatError, Config};
 pub use unbox::unbox;
 
+/// A name of a datatype.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct DatatypeSymbol {
     inner: String,
@@ -43,6 +45,10 @@ impl DatatypeSymbol {
     }
 }
 
+/// A datatype sort.
+///
+/// A datatype sort is a sort that is defined by a datatype. It has a symbol and a list of
+/// argument sorts.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct DatatypeSort {
     symbol: DatatypeSymbol,
@@ -74,6 +80,7 @@ impl DatatypeSort {
     }
 }
 
+/// A sort is the type of a logical term.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub enum Sort {
     Null,
@@ -271,6 +278,10 @@ impl Sort {
 }
 
 rustc_index::newtype_index! {
+    /// An index representing term-level variable.
+    ///
+    /// We manage term-level variables using indices that are unique in each clause.
+    /// [`Clause`] contains `IndexVec<TermVarIdx, Sort>` that manages the indices and the sorts of the variables.
     #[orderable]
     #[debug_format = "v{}"]
     pub struct TermVarIdx { }
@@ -297,6 +308,7 @@ impl TermVarIdx {
     }
 }
 
+/// A known function symbol.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Function {
     name: &'static str,
@@ -376,6 +388,7 @@ impl Function {
     pub const NEG: Function = Function::new("-");
 }
 
+/// A logical term.
 #[derive(Debug, Clone)]
 pub enum Term<V = TermVarIdx> {
     Null,
@@ -693,6 +706,11 @@ impl<V> Term<V> {
 }
 
 rustc_index::newtype_index! {
+    /// An index representing predicate variables.
+    ///
+    /// We manage predicate variables using indices that are unique in a CHC system.
+    /// [`System`] contains `IndexVec<PredVarId, PredVarDef>` that manages the indices
+    /// and signatures of the predicate variables.
     #[debug_format = "p{}"]
     pub struct PredVarId { }
 }
@@ -720,6 +738,9 @@ impl PredVarId {
     }
 }
 
+/// A known predicate.
+///
+/// A known predicate is a predicate that has a fixed meaning, such as `true`, `false`, `=`, etc.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct KnownPred {
     name: &'static str,
@@ -804,6 +825,25 @@ impl KnownPred {
     pub const GREATER_THAN: KnownPred = KnownPred::infix(">");
 }
 
+/// A matcher predicate.
+///
+/// A matcher predicate is a predicate that relates a value of a datatype with its contents.
+/// For example, given the following `enum` datatype:
+///
+/// ```rust
+/// pub enum X {
+///     A(i64),
+///     B(bool),
+/// }
+/// ```
+///
+/// The corresponding matcher predicate is defined as:
+///
+/// ```smtlib2
+/// (define-fun matcher_pred<X> ((x0 Int) (x1 Bool) (v X)) Bool (or (= v (X.A x0)) (= v (X.B x1))))
+/// ```
+///
+/// See the implementation in the [`smtlib2`] module for the details.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct MatcherPred {
     datatype_symbol: DatatypeSymbol,
@@ -851,6 +891,7 @@ impl MatcherPred {
     }
 }
 
+/// A predicate.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Pred {
     Known(KnownPred),
@@ -942,6 +983,7 @@ impl Pred {
     }
 }
 
+/// An atom is a predicate applied to a list of terms.
 #[derive(Debug, Clone)]
 pub struct Atom<V = TermVarIdx> {
     pub pred: Pred,
@@ -1030,6 +1072,11 @@ impl<V> Atom<V> {
     }
 }
 
+/// An arbitrary formula built with atoms and logical connectives.
+///
+/// While it allows arbitrary [`Atom`] in its `Atom` variant, we only expect atoms with known
+/// predicates (i.e., predicates other than `Pred::Var`) to appear in formulas. It is our TODO to
+/// enforce this restriction statically. Also see the definition of [`Body`].
 #[derive(Debug, Clone)]
 pub enum Formula<V = TermVarIdx> {
     Atom(Atom<V>),
@@ -1248,10 +1295,11 @@ impl<V> Formula<V> {
     }
 }
 
+/// The body part of a clause, consisting of atoms and a formula.
 #[derive(Debug, Clone)]
 pub struct Body<V = TermVarIdx> {
     pub atoms: Vec<Atom<V>>,
-    // `formula` doesn't contain PredVar
+    /// NOTE: This doesn't contain predicate variables. Also see [`Formula`].
     pub formula: Formula<V>,
 }
 
@@ -1391,6 +1439,10 @@ where
     }
 }
 
+/// A constrained Horn clause.
+///
+/// A constrained Horn clause is a formula of the form `∀vars. body ⇒ head` where `body` is a conjunction of
+/// atoms and underlying logical formula, and `head` is an atom.
 #[derive(Debug, Clone)]
 pub struct Clause {
     pub vars: IndexVec<TermVarIdx, Sort>,
@@ -1441,19 +1493,26 @@ impl Clause {
     }
 }
 
+/// A selector for a datatype constructor.
+///
+/// A selector is a function that extracts a field from a datatype value.
+/// Through currently we don't use selectors to access datatype fields in [`Term`]s,
+/// we require the symbol as selector name to emit SMT-LIB2 representation of datatypes.
 #[derive(Debug, Clone)]
 pub struct DatatypeSelector {
     pub symbol: DatatypeSymbol,
     pub sort: Sort,
 }
 
+/// A datatype constructor.
 #[derive(Debug, Clone)]
 pub struct DatatypeCtor {
     pub symbol: DatatypeSymbol,
     pub selectors: Vec<DatatypeSelector>,
     pub discriminant: u32,
 }
 
+/// A datatype definition.
 #[derive(Debug, Clone)]
 pub struct Datatype {
     pub symbol: DatatypeSymbol,
@@ -1462,18 +1521,25 @@ pub struct Datatype {
 }
 
 rustc_index::newtype_index! {
+    /// An index of [`Clause`].
+    ///
+    /// [`System`] contains `IndexVec<ClauseId, Clause>` that manages the indices and the clauses.
     #[debug_format = "c{}"]
     pub struct ClauseId { }
 }
 
 pub type PredSig = Vec<Sort>;
 
+/// A predicate variable definition.
 #[derive(Debug, Clone)]
 pub struct PredVarDef {
     pub sig: PredSig,
+    /// We may attach a debug information to include in the resulting SMT-LIB2 file to indicate the
+    /// origin of predicate variables.
     pub debug_info: DebugInfo,
 }
 
+/// A CHC system.
 #[derive(Debug, Clone, Default)]
 pub struct System {
     pub datatypes: Vec<Datatype>,
@@ -1498,6 +1564,13 @@ impl System {
         smtlib2::System::new(self)
     }
 
+    /// Solves the CHC using an external SMT solver.
+    ///
+    /// This method first performs some optimization of the CHC,
+    /// then formats it to SMT-LIB2, and finally calls the configured CHC solver.
+    /// The solver and its arguments can be configured using environment
+    /// variables
+    /// (see <https://github.com/coord-e/thrust?tab=readme-ov-file#configuration>).
     pub fn solve(&self) -> Result<(), CheckSatError> {
         let system = unbox(self.clone());
         if let Ok(file) = std::env::var("THRUST_PRETTY_OUTPUT") {

src/chc/clause_builder.rs

@@ -1,3 +1,11 @@
+//! A builder for [`Clause`]s.
+//!
+//! This module provides [`ClauseBuilder`], a helper for constructing [`Clause`]s. It is
+//! particularly useful for managing the universally quantified variables of a clause. It can
+//! automatically create fresh [`TermVarIdx`] for variables from other domains (e.g.,
+//! [`crate::rty::FunctionParamIdx`]), simplifying the generation of clauses from higher-level
+//! representations.
+
 use std::any::{Any, TypeId};
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -8,6 +16,7 @@ use rustc_index::IndexVec;
 
 use super::{Atom, Body, Clause, DebugInfo, Sort, TermVarIdx};
 
+/// A convenience trait to represent constraints on variables used in [`ClauseBuilder`] at once.
 pub trait Var: Eq + Ord + Hash + Copy + Debug + 'static {}
 impl<T: Eq + Ord + Hash + Copy + Debug + 'static> Var for T {}
 
@@ -54,6 +63,17 @@ impl Hash for dyn Key {
     }
 }
 
+/// A builder for a [`Clause`].
+///
+/// [`Clause`] contains a list of universally quantified variables, a head atom, and a body formula.
+/// When building the head and body, we usually have some formulas that represents variables using
+/// something other than [`TermVarIdx`] (e.g. [`crate::rty::FunctionParamIdx`] or [`crate::refine::Var`]).
+/// These variables are usually OK to be universally quantified in the clause, so we want to keep
+/// the mapping of them to [`TermVarIdx`] and use it to convert the variables in the formulas
+/// during the construction of the clause.
+///
+/// Also see [`crate::rty::ClauseBuilderExt`], which provides a higher-level API on top of this
+/// to build clauses from [`crate::rty::Refinement`]s.
 #[derive(Clone, Default)]
 pub struct ClauseBuilder {
     vars: IndexVec<TermVarIdx, Sort>,

src/chc/debug.rs

@@ -1,3 +1,10 @@
+//! Attachable debug information for CHC clauses.
+//!
+//! The [`DebugInfo`] struct captures contextual information (like `tracing` spans) at the time
+//! of a clause's creation. This information is then pretty-printed as comments in the
+//! generated SMT-LIB2 file, which helps in tracing a clause back to its origin in the
+//! Thrust codebase.
+
 #[derive(Debug, Clone)]
 pub struct Display<'a> {
     inner: &'a DebugInfo,
@@ -19,6 +26,7 @@ impl<'a> std::fmt::Display for Display<'a> {
     }
 }
 
+/// A purely informational metadata that can be attached to a clause.
 #[derive(Debug, Clone, Default)]
 pub struct DebugInfo {
     contexts: Vec<(String, String)>,

src/chc/format_context.rs

@@ -1,7 +1,22 @@
+//! A context for formatting a CHC system into SMT-LIB2.
+//!
+//! This module provides [`FormatContext`], which is responsible for translating parts of [`chc::System`]
+//! into a representation that is compatible with SMT solvers. It handles tasks like
+//! monomorphization of polymorphic datatypes and applying solver-specific workarounds.
+//! The [`super::smtlib2`] module uses this context to perform the final rendering to the SMT-LIB2 format.
+
 use std::collections::BTreeSet;
 
 use crate::chc::{self, hoice::HoiceDatatypeRenamer};
 
+/// A context for formatting a CHC system.
+///
+/// This subsumes a representational difference between [`chc::System`] and resulting SMT-LIB2.
+/// - Gives a naming convention of symbols to represent built-in datatypes of [`chc::System`] in SMT-LIB2,
+/// - Gives a stringified representation of [`chc::Sort`]s,
+/// - Monomorphizes polymorphic datatypes of [`chc::System`] to be compatible with several CHC solvers,
+/// - Renames datatypes to be compatible with Hoice (see [`HoiceDatatypeRenamer`]),
+/// - etc.
 #[derive(Debug, Clone)]
 pub struct FormatContext {
     renamer: HoiceDatatypeRenamer,