Documentation for the compiler (#166)

* getting rid of CompileState.union

* document projections

* more documentation

* fix mdata change; drop Expr.lty

* more docs, some renamings

* more WIP documentation

* more docstrings

* finished Compiler docstrings

* small adjustments

* add `cmpDef` docs

Co-authored-by: Hanting Zhang <hantingz@usc.edu>

TestsUtils/CompileAndExtractTests.lean

@@ -87,8 +87,7 @@ def pairConstants (x y : Array Const) :
 def reindexExpr (map : NatNatMap) : Expr → Expr
   | e@(.var ..)
   | e@(.sort _)
-  | e@(.lit ..)
-  | e@(.lty ..) => e
+  | e@(.lit ..) => e
   | .const n i ls => .const n (map.find! i) ls
   | .app e₁ e₂ => .app (reindexExpr map e₁) (reindexExpr map e₂)
   | .lam n bi e₁ e₂ => .lam n bi (reindexExpr map e₁) (reindexExpr map e₂)

Yatima/Cli/CompileCmd.lean

@@ -18,25 +18,16 @@ def compileRun (p : Cli.Parsed) : IO UInt32 := do
       if !(p.hasFlag "prelude") then setLibsPaths
       let mut stt : CompileState := default
       let log := p.hasFlag "log"
-      let mut errMsg : Option String := none
       let mut cronos := Cronos.new
       for arg in args do
         for filePath in ← getLeanFilePathsList ⟨arg⟩ do
           let filePathStr := filePath.toString
           cronos ← cronos.clock filePathStr
           match ← compile filePath log stt with
-          | .ok stt' => match stt.union stt' with
-            | .ok stt' =>
-              stt := stt'
-              cronos ← cronos.clock filePathStr
-            | .error msg => errMsg := some msg; break
-          | .error msg => errMsg := some (toString msg); break
-        if errMsg.isSome then break
-      match errMsg with
-      | some msg =>
-        IO.eprintln msg
-        return 1
-      | none => pure ()
+          | .ok stt' =>
+            stt := stt'
+            cronos ← cronos.clock filePathStr
+          | .error msg => IO.eprintln msg; return 1
       if p.hasFlag "summary" then
         IO.println s!"{stt.summary}"
         IO.println s!"\n{cronos.summary}"

Yatima/Cli/PipeCmd.lean

@@ -23,29 +23,21 @@ def pipeRun (p : Cli.Parsed) : IO UInt32 := do
       cronos ← cronos.clock "Compilation"
       if !(p.hasFlag "prelude") then setLibsPaths
       let mut stt : CompileState := default
-      let mut errMsg : Option String := none
       let log := p.hasFlag "log"
       let mut cronos' := Cronos.new
       for arg in args do
         for filePath in ← getLeanFilePathsList ⟨arg⟩ do
           let filePathStr := filePath.toString
           cronos' ← cronos'.clock filePathStr
           match ← compile filePath log stt with
-          | .ok stt' => match stt.union stt' with
-            | .ok stt' =>
-              stt := stt'
-              cronos' ← cronos'.clock filePathStr
-            | .error msg => errMsg := some msg; break
-          | .error msg => errMsg := some (toString msg); break
-        if errMsg.isSome then break
-      match errMsg with
-      | some msg =>
-        IO.eprintln msg
-        return 1
-      | none =>
-        if p.hasFlag "summary" then
-          IO.println s!"{stt.summary}"
-          IO.println s!"\n{cronos'.summary}"
+          | .ok stt' =>
+            stt := stt'
+            cronos' ← cronos'.clock filePathStr
+          | .error msg => IO.eprintln msg; return 1
+      if p.hasFlag "summary" then
+        IO.println s!"{stt.summary}"
+        IO.println s!"\n{cronos'.summary}"
+
       cronos ← cronos.clock "Compilation"
       if p.hasFlag "typecheck" then
         cronos ← cronos.clock "Typechecking"

Yatima/Compiler/CompileError.lean

@@ -4,9 +4,11 @@ import Yatima.Datatypes.Const
 
 namespace Yatima.Compiler
 
+/-- Errors that can be thrown in `Yatima.Compiler.CompileM` -/
 inductive CompileError
   | notFoundInCache : Name → CompileError
-  | invalidDereferringIndex : Nat → Nat → CompileError
+  | notFoundInRecrCtx : Name → CompileError
+  | invalidConstantIndex : Nat → Nat → CompileError
   | unknownConstant : Name → CompileError
   | unfilledLevelMetavariable : Lean.Level → CompileError
   | unfilledExprMetavariable : Lean.Expr → CompileError
@@ -26,7 +28,8 @@ inductive CompileError
 
 instance : ToString CompileError where toString
   | .notFoundInCache n => s!"Could not find cid of '{n}' in cache"
-  | .invalidDereferringIndex idx size =>
+  | .notFoundInRecrCtx n => s!"Could not find '{n}' in recrCtx"
+  | .invalidConstantIndex idx size =>
     s!"Invalid index {idx} for dereferring a constant. Must be < {size}."
   | .unknownConstant n => s!"Unknown constant '{n}'"
   | .unfilledLevelMetavariable l => s!"Unfilled level metavariable on universe '{l}'"

Yatima/Compiler/CompileM.lean

@@ -10,30 +10,18 @@ open Std (RBMap)
 /--
 The state for the `Yatima.Compiler.CompileM` monad.
 
-IMPORTANT: `consts` is unreliable when compiling multiple files!
+* `store` contains the resulting set of objects in the IPLD format
+* `consts` is the "pure" array of constants, without CIDs
+* `cache` is just for optimization purposes
 -/
 structure CompileState where
   store  : Ipld.Store
-  cache  : RBMap Name (ConstCid × ConstIdx) compare
   consts : Array Const
+  cache  : RBMap Name (ConstCid × ConstIdx) compare
   deriving Inhabited
 
-namespace CompileState
-
-def union (s s' : CompileState) : Except String CompileState := Id.run do
-  let mut cache := s.cache
-  for (n, c') in s'.cache do
-    match s.cache.find? n with
-    | some c₁ =>
-      if c₁.1 != c'.1 then return throw s!"Conflicting declarations for '{n}'"
-    | none => cache := cache.insert n c'
-  return .ok ⟨
-    s.store.union s'.store,
-    cache,
-    s'.consts
-  ⟩
-
-def summary (s : CompileState) : String :=
+/-- Creates a summary off of a `Yatima.Compiler.CompileState` as a `String` -/
+def CompileState.summary (s : CompileState) : String :=
   let consts := ", ".intercalate $ s.consts.toList.map
     fun c => s!"{c.name} : {c.ctorName}"
   "Compilation summary:\n" ++
@@ -48,77 +36,86 @@ def summary (s : CompileState) : String :=
   s!"  const_meta size: {s.store.const_meta.size}\n" ++
   s!"  cache      size: {s.cache.size}"
 
-end CompileState
+/--
+The type of entries for the `recrCtx`. It contains:
+1. The index of the constant in the mutual block
+2. The index in the list of weakly equal mutual definitions (N/A inductives)
+3. The constant index in array of constants
+-/
+abbrev RecrCtxEntry := (Nat × Option Nat × ConstIdx)
 
-abbrev RecrCtxEntry := (Nat × Option Nat × Nat)
+/--
+The read-only environment for the `Yatima.Compiler.CompileM` monad.
 
+* `constMap` is the original set of constants provided by Lean
+* `univCtx` is the current list of universes
+* `bindCtx` is the current list of binders
+* `recrCtx` is keeps the information for names that represent recursive calls
+* `log` tells whether the user wants to log the compilation
+-/
 structure CompileEnv where
   constMap : Lean.ConstMap
-  univCtx  : List Lean.Name
+  univCtx  : List Name
   bindCtx  : List Name
-  -- (mutual index in recrCtx, weakly equal index (N/A inductives), constant index in array of constants)
   recrCtx  : Std.RBMap Lean.Name RecrCtxEntry compare
   log      : Bool
   deriving Inhabited
 
+/-- Instantiates a `Yatima.Compiler.CompileEnv` from a map of constants -/
 def CompileEnv.init (map : Lean.ConstMap) (log : Bool) : CompileEnv :=
   ⟨map, [], [], .empty, log⟩
 
+/--
+The monad in which compilation takes place is a stack of `ReaderT`, `ExceptT`
+and `StateT` on top of IO
+-/
 abbrev CompileM := ReaderT CompileEnv $
   ExceptT CompileError $ StateT CompileState IO
 
+/-- Basic runner function for `Yatima.Compiler.CompileEnv` -/
 def CompileM.run (env : CompileEnv) (ste : CompileState) (m : CompileM α) :
     IO $ Except CompileError CompileState := do
   match ← StateT.run (ReaderT.run m env) ste with
   | (.ok _,  ste) => return .ok ste
   | (.error e, _) => return .error e
 
-def withName (name : Name) : CompileM α → CompileM α :=
+/-- Computes with a new binder in the monad environment -/
+def withBinder (name : Name) : CompileM α → CompileM α :=
   withReader $ fun e =>
     ⟨e.constMap, e.univCtx, name :: e.bindCtx, e.recrCtx, e.log⟩
 
-def withResetCompileEnv (levels : List Lean.Name) :
+/-- Computes with a given list of levels and reset binders and `recrCtx` -/
+def withLevelsAndReset (levels : List Name) :
     CompileM α → CompileM α :=
   withReader $ fun e => ⟨e.constMap, levels, [], .empty, e.log⟩
 
-def withRecrs (recrCtx : RBMap Lean.Name RecrCtxEntry compare) :
+/-- Computes with a given `recrCtx` -/
+def withRecrs (recrCtx : RBMap Name RecrCtxEntry compare) :
     CompileM α → CompileM α :=
   withReader $ fun e => ⟨e.constMap, e.univCtx, e.bindCtx, recrCtx, e.log⟩
 
-def withLevels (lvls : List Lean.Name) : CompileM α → CompileM α :=
+/-- Computes with a given list of levels-/
+def withLevels (lvls : List Name) : CompileM α → CompileM α :=
   withReader $ fun e => ⟨e.constMap, lvls, e.bindCtx, e.recrCtx, e.log⟩
 
-inductive StoreKey : Type → Type
-  | univ  : Ipld.Both Ipld.UnivCid  → StoreKey (Ipld.Both Ipld.Univ)
-  | expr  : Ipld.Both Ipld.ExprCid  → StoreKey (Ipld.Both Ipld.Expr)
-  | const : Ipld.Both Ipld.ConstCid → StoreKey (Ipld.Both Ipld.Const)
-
-def StoreKey.find? (key : StoreKey A) : CompileM (Option A) := do
-  let store := (← get).store
-  match key with
-  | .univ univCid =>
-    match store.univ_anon.find? univCid.anon, store.univ_meta.find? univCid.meta with
-    | some univAnon, some univMeta => pure $ some ⟨ univAnon, univMeta ⟩
-    | _, _ => pure none
-  | .expr exprCid =>
-    match store.expr_anon.find? exprCid.anon, store.expr_meta.find? exprCid.meta with
-    | some exprAnon, some exprMeta => pure $ some ⟨ exprAnon, exprMeta ⟩
-    | _, _ => pure none
-  | .const constCid =>
-    match store.const_anon.find? constCid.anon, store.const_meta.find? constCid.meta with
-    | some constAnon, some constMeta => pure $ some ⟨ constAnon, constMeta ⟩
-    | _, _ => pure none
-
-def StoreKey.find! (key : StoreKey A) : CompileM A := do
-  let some value ← StoreKey.find? key | throw $ .custom "Cannot find key in store"
-  return value
-
-inductive StoreValue : Type → Type
-  | univ  : Ipld.Both Ipld.Univ  → StoreValue (Ipld.Both Ipld.UnivCid)
-  | expr  : Ipld.Both Ipld.Expr  → StoreValue (Ipld.Both Ipld.ExprCid)
-  | const : Ipld.Both Ipld.Const → StoreValue (Ipld.Both Ipld.ConstCid)
-
-def StoreValue.insert : StoreValue A → CompileM A
+/-- Possibly gets a `Yatima.Compiler.RecrCtxEntry` from the `recrCtx` by name -/
+def getFromRecrCtx (name : Name) : CompileM $ Option RecrCtxEntry :=
+  return (← read).recrCtx.find? name
+
+/-- Forcibly gets a `Yatima.Compiler.RecrCtxEntry` from the `recrCtx` by name -/
+def getFromRecrCtx! (name : Name) : CompileM $ RecrCtxEntry := do
+  match ← getFromRecrCtx name with
+  | some entry => pure entry
+  | none => throw $ .notFoundInRecrCtx name
+
+/-- Auxiliary type to standardize additions of CIDs to the store -/
+inductive StoreEntry : Type → Type
+  | univ  : Ipld.Both Ipld.Univ  → StoreEntry (Ipld.Both Ipld.UnivCid)
+  | expr  : Ipld.Both Ipld.Expr  → StoreEntry (Ipld.Both Ipld.ExprCid)
+  | const : Ipld.Both Ipld.Const → StoreEntry (Ipld.Both Ipld.ConstCid)
+
+/-- Adds CID data to the store, but also returns it for practical reasons -/
+def addToStore : StoreEntry A → CompileM A
   | .univ  obj =>
     let cid := ⟨ ToIpld.univToCid obj.anon, ToIpld.univToCid obj.meta ⟩
     modifyGet (fun stt => (cid, { stt with store :=
@@ -132,7 +129,7 @@ def StoreValue.insert : StoreValue A → CompileM A
   | .const obj =>
     let cid := ⟨ ToIpld.constToCid obj.anon, ToIpld.constToCid obj.meta ⟩
     match obj.anon, obj.meta with
-    -- Mutual definition/inductive blocks do not get added to the set of definitions
+    -- Mutual definition/inductive blocks do not get added to the set of constants
     | .mutDefBlock .., .mutDefBlock ..
     | .mutIndBlock .., .mutIndBlock .. =>
       modifyGet fun stt => (cid, { stt with store :=
@@ -144,14 +141,16 @@ def StoreValue.insert : StoreValue A → CompileM A
                          const_meta := stt.store.const_meta.insert cid.meta obj.meta,
                          consts     := stt.store.consts.insert cid } })
 
+/-- Adds data associated with a name to the cache -/
 def addToCache (name : Name) (c : ConstCid × ConstIdx) : CompileM Unit := do
   modify fun stt => { stt with cache := stt.cache.insert name c }
 
-def addToConsts (idx : Nat) (c : Const): CompileM Unit := do
+/-- Adds a constant to the array of constants at a given index -/
+def addToConsts (idx : ConstIdx) (c : Const) : CompileM Unit := do
   let consts := (← get).consts
   if h : idx < consts.size then
     modify fun stt => { stt with consts := consts.set ⟨idx, h⟩ c }
   else
-    throw $ .invalidDereferringIndex idx consts.size
+    throw $ .invalidConstantIndex idx consts.size
 
 end Yatima.Compiler