diff --git a/Lampe/Lampe/Builtin/Array.lean b/Lampe/Lampe/Builtin/Array.lean index c0107622..c2c6c862 100644 --- a/Lampe/Lampe/Builtin/Array.lean +++ b/Lampe/Lampe/Builtin/Array.lean @@ -91,6 +91,58 @@ def mkRepeatedArray := newGenericTotalPureBuiltin (fun (len, tp) => ⟨[tp], (.array tp len)⟩) (fun (num, _) h![val] => List.Vector.replicate num.toNat val) +/-- +Interprets a list of element values as an array of length `n`, truncating or zero-padding as +needed. An array denotes as a `List.Vector`, i.e. a list *paired with a proof* that its length +is `n`; both components are built below so that they are independent of the concrete elements. +-/ +def valArray {p : Prime} (tp : Tp) (n : U 32) (vals : List (Tp.denote p tp)) : + Tp.denote p (.array tp n) := + ⟨-- Reshape the input to exactly `n` elements: keep the first `n`, pad with the type's zero + -- value if `vals` is too short. The elaborator always supplies a list of exactly the right + -- length, so this is the identity in practice; it exists so the result is length-correct + -- *by construction* for any input. + vals.takeD n.toNat (Tp.zero p tp), + -- The generic lemma `(l.takeD n d).length = n`. Because `takeD` guarantees the length for + -- *any* list, this discharges the vector's length side-condition without ever inspecting + -- the elements — checking a 10'000-element literal costs the same as a 1-element one, + -- unlike proving `[a, b, c, …].length = n` by `rfl`/`decide`, which walks the whole list. + List.takeD_length _ _ _⟩ + +/-- +Defines the builtin constructor for arrays whose elements are all compile-time constants. + +The element values are carried by the builtin itself as a (prime-generic) denoted list, rather +than as per-element expressions. The Lampe elaborator emits this builtin — with `vals` referencing +a hoisted auxiliary definition — for array literals all of whose elements are numeric literals. +This keeps both the extracted term and every proof goal mentioning the array shallow (a single +constant), in contrast to the general `mkArray` path which `letIn`-binds each element and so +produces terms whose depth grows with the array length. +-/ +def mkValArray + (arrTp : Tp) + -- `vals` is a *function of the prime* rather than a plain list, because element values live + -- in `Tp.denote p _`, which depends on `p` — e.g. `Field` elements are integers modulo `p`. + -- The builtin must work at every prime, so the hoisted definition abstracts over it. + (vals : (p : Prime) → List (Tp.denote p arrTp.arrayElem)) := + newGenericTotalPureBuiltin + -- The generic-argument type is `Unit` (unlike `mkArray`, which is indexed by element count + -- and type): the builtin is fully determined by its parameters `arrTp`/`vals`, so proof + -- automation always instantiates it with `a := ()`. The signature takes *no runtime + -- arguments* (`[]`) — the elements are data inside the builtin — so stepping over a call + -- produces no per-element subgoals. The output type is recovered from the whole array type + -- via the projections `arrayElem`/`arraySize` rather than taking element type and length + -- as separate parameters: the elaborator can pass the one type annotation it already has, + -- and because the projections are `@[reducible]`, `(.array tp n).arrayElem`/`.arraySize` + -- unify with `tp`/`n` when the closing lemma from `Tactic/Steps.lean` — stated at the + -- projected type — is matched against a goal stated at the original array type. + (fun (_ : Unit) => ⟨[], .array arrTp.arrayElem arrTp.arraySize⟩) + -- Evaluation: `h![]` matches the empty `HList` of runtime arguments, and the result is + -- `valArray … (vals p)` (the prime is implicit from context), whose length side-condition + -- is discharged generically (see `valArray`) — no proof obligation about the concrete + -- elements ever arises. + (fun _ h![] => valArray arrTp.arrayElem arrTp.arraySize (vals _)) + /-- Defines the indexing of a array `l : Array tp n` with `i : U 32` We make the following assumptions: diff --git a/Lampe/Lampe/Builtin/Vector.lean b/Lampe/Lampe/Builtin/Vector.lean index 9aa9811f..afc4e716 100644 --- a/Lampe/Lampe/Builtin/Vector.lean +++ b/Lampe/Lampe/Builtin/Vector.lean @@ -30,6 +30,19 @@ def mkRepeatedVector := newGenericTotalPureBuiltin (fun (a : Tp)=> ⟨[Tp.u 32, a], (.vector a)⟩) (fun _ h![n, val] => List.replicate n.toNat val) +/-- +Defines the builtin constructor for vectors (slices) whose elements are all compile-time +constants. The element values are carried by the builtin itself as a (prime-generic) denoted +list; see `Builtin.mkValArray` for a step-by-step account of the construction. +-/ +def mkValVector (tp : Tp) (vals : (p : Prime) → List (Tp.denote p tp)) := + newGenericTotalPureBuiltin + (fun (_ : Unit) => ⟨[], .vector tp⟩) + -- Simpler than the array case: a vector denotes as a bare `List` with no length invariant, + -- so `vals p` is already a value of the output type and is returned directly — no `takeD` + -- reshaping and no length lemma are needed. + (fun _ h![] => vals _) + /-- Defines the indexing of a vector `l : List tp` with `i : U 32` We make the following assumptions: diff --git a/Lampe/Lampe/Syntax/Elab.lean b/Lampe/Lampe/Syntax/Elab.lean index ae0334e0..b433e915 100644 --- a/Lampe/Lampe/Syntax/Elab.lean +++ b/Lampe/Lampe/Syntax/Elab.lean @@ -11,18 +11,125 @@ open Lean Elab -- DSL: TERMS ------------------------------------------------------------------------------------- +/-- +Views `stx` as a compile-time-constant element of an array literal, returning a term for its +denoted value. Returns `none` for anything that is not a numeric or boolean literal. + +Note that the element's own Noir-level type ascription (`$n : $_`) is dropped: the returned +term is a bare Lean numeral (or `Bool` literal), which later elaborates against the expected +element type `Tp.denote p tp` imposed by the hoisted definition's type ascription — via that +type's `OfNat` (resp. `Neg`) instance. Dropping the ascription is sound because extraction +always ascribes every element with the array's element type. +-/ +private partial def litArrayElem (stx : TSyntax `noir_expr) : + Elab.Command.CommandElabM (Option (TSyntax `term)) := do + match stx with + | `(noir_expr|($e:noir_expr)) => litArrayElem e + | `(noir_expr|$n:num : $_) => return some (←`($n)) + | `(noir_expr|-$n:num : $_) => return some (←`((-$n))) + | `(noir_expr|#_true) => return some (←`(true)) + | `(noir_expr|#_false) => return some (←`(false)) + | _ => return none + +/-- Applies `f` to every node of `stx` top-down, replacing a node (without descending into the +replacement) whenever `f` returns `some`. -/ +private partial def replaceSyntaxTopDownM [Monad m] (f : Syntax → m (Option Syntax)) + (stx : Syntax) : m Syntax := do + match ← f stx with + | some new => return new + | none => match stx with + | .node info kind args => return .node info kind (← args.mapM (replaceSyntaxTopDownM f)) + | s => return s + +/-- +Rewrites every `#_ mkArray` / `#_ mkVector` call in `stx` whose arguments are all compile-time +constants into a single `Builtin.mkValArray` / `Builtin.mkValVector` call (via the `splice!` +escape hatch), hoisting the element values into an auxiliary definition named +`«#lits»`. + +The general `mkArray` path `letIn`-binds every element, so an `n`-element literal produces a term +(and, later, proof goals) of depth `O(n)`; every recursive traversal of such a term — during +elaboration, `simp`, or unification — then needs recursion depth and time proportional to `n`, +which makes large array literals unusably slow. After this rewrite the extracted body and all +goals about it contain only the (shallow) auxiliary constant, so their cost is independent of the +array size. The auxiliary definition elaborates the deep list literal exactly once. + +Arrays with any non-constant element (e.g. a function call) are left on the general path. +-/ +private def hoistLiteralArrays (baseName : Name) (stx : Syntax) : + Elab.Command.CommandElabM Syntax := do + let counter ← IO.mkRef (0 : Nat) + let pId := mkIdent `p + let rewrite (node : Syntax) : Elab.Command.CommandElabM (Option Syntax) := do + let tnode : TSyntax `noir_expr := ⟨node⟩ + match tnode with + | `(noir_expr|(#_ $nm:ident returning $tp)( $args,* )) => do + let isArray := nm.getId == `mkArray + unless isArray || nm.getId == `mkVector do return none + -- Empty literals stay on the general path: they are already cheap, and there would be + -- nothing to hoist. + if args.getElems.isEmpty then return none + -- All-or-nothing: view every element as a compile-time constant; if any element is not + -- one (`mapM id` folds `Array (Option _)` into `Option (Array _)`), leave the whole + -- literal on the general `mkArray`/`mkVector` path. + let elems ← args.getElems.mapM litArrayElem + let some elems := elems.mapM id | return none + -- Elaborate the surface type annotation into a `Tp` term; it becomes the builtin's + -- parameter and, projected via `Tp.arrayElem`/`Tp.vectorElem`, the element type of the + -- hoisted list. + let arrTp ← MonadDSL.run (makeNoirType tp) + -- A `[…]` literal expands to a right-nested chain of `List.cons`, so a single literal of + -- all `n` elements would make the elaborator recurse to depth `n` (overflowing its stack + -- for large `n`). Emitting the elements as 256-element `[…]` chunks joined by `++` caps + -- the depth contributed by any one literal at 256, leaving only the far shallower + -- (`n / 256`-deep) `++` spine. + let chunkSize := 256 + let mut chunks : Array (Array (TSyntax `term)) := #[] + let mut i := 0 + while i < elems.size do + chunks := chunks.push (elems.extract i (min (i + chunkSize) elems.size)) + i := i + chunkSize + let mut body ← `([$(chunks.back!),*]) + for c in chunks.pop.reverse do + body ← `([$c,*] ++ $body) + let idx ← counter.modifyGet fun i => (i, i + 1) + -- `#` cannot appear in Noir identifiers, so the auxiliary name can never collide with an + -- extracted definition. + let auxId := mkIdent <| Name.mkSimple s!"{baseName.getString!}#lits{idx}" + let elemTp ← if isArray then `(Tp.arrayElem $arrTp) else `(Tp.vectorElem $arrTp) + -- Elaborate the auxiliary definition immediately, as its own command: the deep element + -- list is type-checked exactly once, here, and every later mention of the array — in the + -- extracted body and in proof goals — is just this constant. + Elab.Command.elabCommand <| ← + `(def $auxId ($pId : Prime) : List (Tp.denote $pId $elemTp) := $body) + -- Replace the literal with a raw `Expr.callBuiltin` via the `splice!` escape hatch + -- (bypassing the surface grammar, which has no syntax for data-carrying builtins). The + -- call takes no argument expressions — the elements travel inside the builtin. + let repl ← if isArray then + `(noir_expr| + splice!( Expr.callBuiltin [] $arrTp (Builtin.mkValArray $arrTp $auxId) h![] )) + else + `(noir_expr| + splice!( Expr.callBuiltin [] $arrTp (Builtin.mkValVector $elemTp $auxId) h![] )) + return some repl.raw + | _ => return none + replaceSyntaxTopDownM rewrite stx + /-- Elaborates a function definition written in the Noir eDSL. -/ elab d:noir_depr? "noir_def" decl:noir_fn_def : command => do + let baseName ← makeNoirIdent decl.raw[0] + let decl : TSyntax `noir_fn_def := ⟨← hoistLiteralArrays baseName.getId decl.raw⟩ let (name, decl) ← makeFnDecl decl let decl ← match (←parseDeprecatedMessage d) with | some msg => `( - @[deprecated $name $(Syntax.mkStrLit msg) (since := "")] + @[deprecated $name $(Syntax.mkStrLit msg) (since := "")] def $name : FunctionDecl := $decl) | none => `(def $name : FunctionDecl := $decl) Elab.Command.elabCommand decl /-- Elaborates a trait implementation written in the Noir eDSL. -/ elab "noir_trait_impl[" defName:ident "]" impl:noir_trait_impl : command => do + let impl : TSyntax `noir_trait_impl := ⟨← hoistLiteralArrays defName.getId impl.raw⟩ let (name, impl) ← makeTraitImpl impl let decl ← `(def $defName : String × TraitImpl := ($(Syntax.mkStrLit name.getId.toString), $impl)) Elab.Command.elabCommand decl diff --git a/Lampe/Lampe/Tactic/Steps.lean b/Lampe/Lampe/Tactic/Steps.lean index 8cd1584e..99ee5ef3 100644 --- a/Lampe/Lampe/Tactic/Steps.lean +++ b/Lampe/Lampe/Tactic/Steps.lean @@ -158,6 +158,23 @@ def getClosingTerm (val : Lean.Expr) : TacticM (Option (TSyntax `term)) := withT return some (←``(genericTotalPureBuiltin_intro Builtin.mkArray (a := (List.length $argTypes, _)) rfl)) | ``Lampe.Builtin.mkRepeatedArray => return some (←``(genericTotalPureBuiltin_intro Builtin.mkRepeatedArray (a := (_, _)) rfl)) + | ``Lampe.Builtin.mkValArray => + -- `mkValArray` is a parameterized builtin, so the goal mentions it fully applied as + -- `Builtin.mkValArray arrTp vals`. Pull those two arguments out of the goal's + -- builtin expression… + let some arrTp := builtin.getAppArgs[0]? | throwError "malformed mkValArray" + let some vals := builtin.getAppArgs[1]? | throwError "malformed mkValArray" + -- …reify them back into syntax… + let arrTp ← arrTp.toSyntax + let vals ← vals.toSyntax + -- …and re-apply them inside the closing term, so that the intro lemma's hypothesis + -- `b = newGenericTotalPureBuiltin sgn desc` unfolds `mkValArray` one step and closes + -- by `rfl`, inferring `sgn`/`desc` from that equation. The generic argument is `()` + -- since `mkValArray` carries all its data in its parameters. The resulting + -- postcondition is `v = Builtin.valArray _ _ (vals p)` — a single shallow constant, + -- regardless of how many elements the literal has. + return some + (←``(genericTotalPureBuiltin_intro (Builtin.mkValArray $arrTp $vals) (a := ()) rfl)) | ``Lampe.Builtin.arrayIndex => return some (←``(arrayIndex_intro)) | ``Lampe.Builtin.arrayLen => let some argTps := val.getAppArgs[1]? | throwError "malformed arrayLen" @@ -176,6 +193,15 @@ def getClosingTerm (val : Lean.Expr) : TacticM (Option (TSyntax `term)) := withT return some (←``(genericTotalPureBuiltin_intro Builtin.mkVector (a := (List.length $argTypes, _)) rfl)) | ``Lampe.Builtin.mkRepeatedVector => return some (←``(genericTotalPureBuiltin_intro Builtin.mkRepeatedVector (a := _) rfl)) + | ``Lampe.Builtin.mkValVector => + -- Same shape as the `mkValArray` case above, with the element type in place of the + -- whole array type. + let some tp := builtin.getAppArgs[0]? | throwError "malformed mkValVector" + let some vals := builtin.getAppArgs[1]? | throwError "malformed mkValVector" + let tp ← tp.toSyntax + let vals ← vals.toSyntax + return some + (←``(genericTotalPureBuiltin_intro (Builtin.mkValVector $tp $vals) (a := ()) rfl)) | ``Lampe.Builtin.vectorPushBack => return some (←``(genericTotalPureBuiltin_intro Builtin.vectorPushBack rfl)) | ``Lampe.Builtin.vectorPushFront => return some (←``(genericTotalPureBuiltin_intro Builtin.vectorPushFront rfl)) | ``Lampe.Builtin.vectorIndex => return some (←``(vectorIndex_intro)) diff --git a/Lampe/Lampe/Tp.lean b/Lampe/Lampe/Tp.lean index 95a067ae..3d111a70 100644 --- a/Lampe/Lampe/Tp.lean +++ b/Lampe/Lampe/Tp.lean @@ -281,6 +281,34 @@ match tp with end +/-- +The element type of an array type (or `.unit` for non-array types). + +This and the projections below exist so that `Builtin.mkValArray`/`mkValVector` can state their +signatures in terms of the *whole* array/vector type they are given. They are `@[reducible]` so +that e.g. `(Tp.array tp n).arrayElem` reduces to `tp` during unification — the Hoare triple +produced for a `mkValArray` call (see `Tactic/Steps.lean`) is stated at the projected type +`.array arrTp.arrayElem arrTp.arraySize`, and must unify with goals stated at the original +array type. The `.unit`/`0` fallbacks for non-matching types are never reached by the +elaborator; they just make the projections total. +-/ +@[reducible] +def Tp.arrayElem : Tp → Tp + | .array tp _ => tp + | _ => .unit + +/-- The length of an array type (or `0` for non-array types). -/ +@[reducible] +def Tp.arraySize : Tp → U 32 + | .array _ n => n + | _ => 0 + +/-- The element type of a vector (slice) type (or `.unit` for non-vector types). -/ +@[reducible] +def Tp.vectorElem : Tp → Tp + | .vector tp => tp + | _ => .unit + /-- Index into a `Tp.denoteArgs` tuple by `Member` witness. -/ @[reducible] def Tp.denoteArgs.getByMember : {tps : List Tp} → Tp.denoteArgs p tps → Member tp tps → Tp.denote p tp