U8 type (#426)

* `U8` type. Type safe `u8` operations

* `U8` pattern matching

* `U8` in the IxVM

* `Sha256` fix

Author: gabriel-barrett

Ix/Aiur/Compiler/Check.lean

@@ -48,6 +48,7 @@ inductive CheckError
   | unconstrainedConstructor : Global → CheckError
   | infiniteType : Nat → Typ → CheckError
   | unresolvedMVar : Nat → CheckError
+  | u8LitOutOfRange : Nat → CheckError
   deriving Repr
 
 instance : ToString CheckError where
@@ -59,6 +60,7 @@ open Source
 def Typ.instantiate (subst : Global → Option Typ) : Typ → Typ
   | .unit => .unit
   | .field => .field
+  | .u8 => .u8
   | .tuple ts => .tuple (ts.attach.map (fun ⟨t, _⟩ => Typ.instantiate subst t))
   | .array t n => .array (Typ.instantiate subst t) n
   | .pointer t => .pointer (Typ.instantiate subst t)
@@ -83,6 +85,7 @@ def expandTypeMBound : Nat → Std.HashSet Global → Array TypeAlias →
     Typ → StateT (Std.HashMap Global Typ) (Except CheckError) Typ
   | _, _, _, .unit => pure .unit
   | _, _, _, .field => pure .field
+  | _, _, _, .u8 => pure .u8
   | bound, visited, tops, .pointer t => do
     pure $ .pointer (← expandTypeMBound bound visited tops t)
   | bound, visited, tops, .function inputs output => do
@@ -307,7 +310,7 @@ def unifyTypBound : Nat → Typ → Typ → CheckM Bool
       if a == b then pure true else do bindMVar a (.mvar b); pure true
     | .mvar a, _ => do bindMVar a t2; pure true
     | _, .mvar b => do bindMVar b t1; pure true
-    | .unit, .unit | .field, .field => pure true
+    | .unit, .unit | .field, .field | .u8, .u8 => pure true
     | .tuple ts1, .tuple ts2 =>
       if ts1.size != ts2.size then pure false else
         ts1.zip ts2 |>.allM fun (x, y) => unifyTypBound bound x y
@@ -330,6 +333,7 @@ available proxy for `sizeOf` in `unifyTyp`'s bound. -/
 def Typ.nodeCount : Typ → Nat
   | .unit        => 1
   | .field       => 1
+  | .u8          => 1
   | .mvar _      => 1
   | .ref _       => 1
   | .pointer t   => 1 + Typ.nodeCount t
@@ -445,8 +449,10 @@ def checkPatternAux (pat : Pattern) (typ : Typ) : CheckM (List (Local × Typ)) :
   | .var var => pure [(var, typ)]
   | .wildcard => pure []
   | .field _ =>
+    -- A numeric-literal pattern matches both `field` and `u8` scrutinees: a
+    -- `u8` is a field element, so comparing it against a byte constant is sound.
     match typ with
-    | .field => pure []
+    | .field | .u8 => pure []
     | _ => throw $ .incompatiblePattern pat typ
   | .tuple pats =>
     match typ with
@@ -711,54 +717,73 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
     let b' ← checkNoEscape b .field
     pure (Typed.Term.mul .field false a' b')
   | .u8ShiftLeft a => do
-    let a' ← checkNoEscape a .field
-    pure (Typed.Term.u8ShiftLeft .field false a')
+    let a' ← checkNoEscape a .u8
+    pure (Typed.Term.u8ShiftLeft .u8 false a')
   | .u8ShiftRight a => do
-    let a' ← checkNoEscape a .field
-    pure (Typed.Term.u8ShiftRight .field false a')
+    let a' ← checkNoEscape a .u8
+    pure (Typed.Term.u8ShiftRight .u8 false a')
   | .u8BitDecomposition a => do
-    let a' ← checkNoEscape a .field
+    -- Bits are 0/1 values, kept as plain `field`.
+    let a' ← checkNoEscape a .u8
     pure (Typed.Term.u8BitDecomposition (.array .field 8) false a')
   | .u8Xor a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8Xor .field false a' b')
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8Xor .u8 false a' b')
   | .u8And a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8And .field false a' b')
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8And .u8 false a' b')
   | .u8Or a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8Or .field false a' b')
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8Or .u8 false a' b')
   | .u8Add a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8Add (.tuple #[.field, .field]) false a' b')
+    -- Low byte and the 0/1 carry are both `u8` (the carry is provably in range:
+    -- the add lookup forces the inputs to be bytes, so `carry ∈ {0, 1}`).
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8Add (.tuple #[.u8, .u8]) false a' b')
   | .u8Mul a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8Mul (.tuple #[.field, .field]) false a' b')
+    -- Both low and high bytes are `u8`.
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8Mul (.tuple #[.u8, .u8]) false a' b')
   | .u8ChainRotr7 a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8ChainRotr7 (.tuple #[.field, .field, .field]) false a' b')
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8ChainRotr7 (.tuple #[.u8, .u8, .u8]) false a' b')
   | .u8ChainRotr4 a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8ChainRotr4 (.tuple #[.field, .field, .field]) false a' b')
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8ChainRotr4 (.tuple #[.u8, .u8, .u8]) false a' b')
   | .u8Sub a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
-    pure (Typed.Term.u8Sub (.tuple #[.field, .field]) false a' b')
+    -- Low byte and the 0/1 borrow are both `u8` (same range argument as add).
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
+    pure (Typed.Term.u8Sub (.tuple #[.u8, .u8]) false a' b')
   | .u8LessThan a b => do
-    let a' ← checkNoEscape a .field
-    let b' ← checkNoEscape b .field
+    -- Result is a 0/1 flag (`field`).
+    let a' ← checkNoEscape a .u8
+    let b' ← checkNoEscape b .u8
     pure (Typed.Term.u8LessThan .field false a' b')
   | .u32LessThan a b => do
     let a' ← checkNoEscape a .field
     let b' ← checkNoEscape b .field
     pure (Typed.Term.u32LessThan .field false a' b')
+  | .u8Lit n => do
+    if n ≥ 256 then throw (.u8LitOutOfRange n)
+    pure (Typed.Term.field .u8 false (G.ofNat n))
+  | .u8RangeCheck a b => do
+    let a' ← checkNoEscape a .field
+    let b' ← checkNoEscape b .field
+    pure (Typed.Term.u8RangeCheck (.tuple #[.u8, .u8]) false a' b')
+  | .toField a => do
+    let a' ← checkNoEscape a .u8
+    pure (Typed.Term.toField .field false a')
+  | .u8FromFieldUnsafe a => do
+    let a' ← checkNoEscape a .field
+    pure (Typed.Term.u8FromFieldUnsafe .u8 false a')
   | .ioGetInfo key => do
     let key' ← inferNoEscape key
     match ← walkTyp key'.typ with
@@ -922,6 +947,11 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
       pure (.u8LessThan (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
   | .u32LessThan τ e a b => do
       pure (.u32LessThan (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
+  | .u8RangeCheck τ e a b => do
+      pure (.u8RangeCheck (← zonkTyp τ) e (← zonkTypedTerm a) (← zonkTypedTerm b))
+  | .toField τ e a => do pure (.toField (← zonkTyp τ) e (← zonkTypedTerm a))
+  | .u8FromFieldUnsafe τ e a => do
+      pure (.u8FromFieldUnsafe (← zonkTyp τ) e (← zonkTypedTerm a))
   | .debug τ e label t r => do
       let t' ← match t with
         | none => pure none

Ix/Aiur/Compiler/Concretize.lean

@@ -50,6 +50,9 @@ def typToConcrete (mono : Std.HashMap (Global × Array Typ) Global) :
     Typ → Except ConcretizeError Concrete.Typ
   | .unit => pure .unit
   | .field => pure .field
+  -- `u8` is erased here: same representation as `field`, distinction no longer
+  -- needed past type-checking.
+  | .u8 => pure .field
   | .tuple ts => do
       pure (.tuple (← ts.attach.mapM fun ⟨t, _⟩ => typToConcrete mono t))
   | .array t n => do pure (.array (← typToConcrete mono t) n)
@@ -77,6 +80,7 @@ termination_by t => sizeOf t
 
 def Typ.toFlatName : Typ → String
   | .field => "G"
+  | .u8 => "U8"
   | .unit => "Unit"
   | Typ.ref g => g.toName.toString (escape := false)
   | .pointer t => "Ptr_" ++ t.toFlatName
@@ -93,6 +97,7 @@ decreasing_by all_goals first | decreasing_tactic | grind
 
 def Typ.appendNameLimbs (g : Global) : Typ → Global
   | .field => g.pushNamespace "G"
+  | .u8 => g.pushNamespace "U8"
   | .unit => g.pushNamespace "Unit"
   | Typ.ref g' =>
     let rec pushAll (g : Global) : Lean.Name → Global
@@ -349,6 +354,12 @@ def termToConcrete
   | .u32LessThan τ e a b => do
       pure (.u32LessThan (← typToConcrete mono τ) e
                          (← termToConcrete mono a) (← termToConcrete mono b))
+  | .u8RangeCheck τ e a b => do
+      pure (.u8RangeCheck (← typToConcrete mono τ) e
+                          (← termToConcrete mono a) (← termToConcrete mono b))
+  -- `toField` / `u8FromFieldUnsafe` are erased coercions: `u8` and `field`
+  -- share a representation, so we drop the wrapper and keep the inner term.
+  | .toField _ _ a | .u8FromFieldUnsafe _ _ a => termToConcrete mono a
   | .debug τ e l t r => do
       -- Inline the Option.mapM case-split so termination sees the sub-Term.
       let t' ← match t with
@@ -546,6 +557,12 @@ def rewriteTypedTerm (decls : Typed.Decls)
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
   | .u32LessThan τ e a b => .u32LessThan (rewriteTyp subst mono τ) e
       (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
+  | .u8RangeCheck τ e a b => .u8RangeCheck (rewriteTyp subst mono τ) e
+      (rewriteTypedTerm decls subst mono a) (rewriteTypedTerm decls subst mono b)
+  | .toField τ e a => .toField (rewriteTyp subst mono τ) e
+      (rewriteTypedTerm decls subst mono a)
+  | .u8FromFieldUnsafe τ e a => .u8FromFieldUnsafe (rewriteTyp subst mono τ) e
+      (rewriteTypedTerm decls subst mono a)
   | .debug τ e l t r =>
     let t' := match t with
       | none => none
@@ -613,11 +630,12 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
     args.attach.foldl (fun s ⟨a, _⟩ => collectInTypedTerm s a) seen
   | .add τ _ a b | .sub τ _ a b | .mul τ _ a b
   | .u8Xor τ _ a b | .u8Add τ _ a b | .u8Mul τ _ a b | .u8Sub τ _ a b
-  | .u8ChainRotr7 τ _ a b | .u8ChainRotr4 τ _ a b
+  | .u8ChainRotr7 τ _ a b | .u8ChainRotr4 τ _ a b | .u8RangeCheck τ _ a b
   | .u8And τ _ a b | .u8Or τ _ a b
   | .u8LessThan τ _ a b | .u32LessThan τ _ a b =>
     collectInTypedTerm (collectInTypedTerm (collectInTyp seen τ) a) b
-  | .eqZero τ _ a | .store τ _ a | .load τ _ a | .ptrVal τ _ a
+  | .eqZero τ _ a | .store τ _ a | .load τ _ a | .ptrVal τ _ a | .toField τ _ a
+  | .u8FromFieldUnsafe τ _ a
   | .u8BitDecomposition τ _ a | .u8ShiftLeft τ _ a | .u8ShiftRight τ _ a
   | .ioGetInfo τ _ a => collectInTypedTerm (collectInTyp seen τ) a
   | .proj τ _ a _ | .get τ _ a _ | .slice τ _ a _ _ =>
@@ -676,11 +694,12 @@ def collectCalls (decls : Typed.Decls)
     bs.attach.foldl (fun s ⟨(_, b), _⟩ => collectCalls decls s b) seen
   | .add _ _ a b | .sub _ _ a b | .mul _ _ a b
   | .u8Xor _ _ a b | .u8Add _ _ a b | .u8Mul _ _ a b | .u8Sub _ _ a b
-  | .u8ChainRotr7 _ _ a b | .u8ChainRotr4 _ _ a b
+  | .u8ChainRotr7 _ _ a b | .u8ChainRotr4 _ _ a b | .u8RangeCheck _ _ a b
   | .u8And _ _ a b | .u8Or _ _ a b
   | .u8LessThan _ _ a b | .u32LessThan _ _ a b =>
     collectCalls decls (collectCalls decls seen a) b
-  | .eqZero _ _ a | .store _ _ a | .load _ _ a | .ptrVal _ _ a
+  | .eqZero _ _ a | .store _ _ a | .load _ _ a | .ptrVal _ _ a | .toField _ _ a
+  | .u8FromFieldUnsafe _ _ a
   | .u8BitDecomposition _ _ a | .u8ShiftLeft _ _ a | .u8ShiftRight _ _ a
   | .ioGetInfo _ _ a => collectCalls decls seen a
   | .proj _ _ a _ | .get _ _ a _ | .slice _ _ a _ _ => collectCalls decls seen a
@@ -782,6 +801,11 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
       (substInTypedTerm subst a) (substInTypedTerm subst b)
   | .u32LessThan τ e a b => .u32LessThan (Typ.instantiate subst τ) e
       (substInTypedTerm subst a) (substInTypedTerm subst b)
+  | .u8RangeCheck τ e a b => .u8RangeCheck (Typ.instantiate subst τ) e
+      (substInTypedTerm subst a) (substInTypedTerm subst b)
+  | .toField τ e a => .toField (Typ.instantiate subst τ) e (substInTypedTerm subst a)
+  | .u8FromFieldUnsafe τ e a =>
+    .u8FromFieldUnsafe (Typ.instantiate subst τ) e (substInTypedTerm subst a)
   | .debug τ e l t r =>
     let t' := match t with
       | none => none

Ix/Aiur/Compiler/Layout.lean

@@ -198,6 +198,8 @@ def opLayout : Bytecode.Op → LayoutM Unit
   | .u8ChainRotr7 .. | .u8ChainRotr4 .. => do pushDegrees #[1, 1, 1]; bumpAuxiliaries 3; bumpLookups
   | .u8LessThan .. => do pushDegree 1; bumpAuxiliaries; bumpLookups
   | .u32LessThan .. => do pushDegree 1; bumpAuxiliaries 12; bumpLookups 6
+  -- Pure range-check lookup: no output columns/degrees, just one lookup.
+  | .u8RangeCheck .. => bumpLookups
   | .debug .. => pure ()
 
 /-- Termination helper for blockLayout's Block/Ctrl traversal. -/

Ix/Aiur/Compiler/Lower.lean

@@ -299,6 +299,13 @@ def toIndex
   | .u8Or _ _ i j => do let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j; pushOp (.u8Or i j)
   | .u8LessThan _ _ i j => do let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j; pushOp (.u8LessThan i j)
   | .u32LessThan _ _ i j => do let i ← expectIdx layoutMap bindings i; let j ← expectIdx layoutMap bindings j; pushOp (.u32LessThan i j)
+  | .u8RangeCheck _ _ i j => do
+    -- Side-effecting lookup; the two `u8` outputs alias the inputs, so no new
+    -- value slots are allocated (cf. `.assertEq`).
+    let i ← expectIdx layoutMap bindings i
+    let j ← expectIdx layoutMap bindings j
+    modify fun stt => { stt with ops := stt.ops.push (.u8RangeCheck i j) }
+    pure #[i, j]
   | .debug _ _ label term ret => do
     let term ← match term with
       | none => pure none

Ix/Aiur/Compiler/Match.lean

@@ -388,6 +388,9 @@ def typedToSimple : Term → Simple.Term
   | .u8Or τ e a b => .u8Or τ e (typedToSimple a) (typedToSimple b)
   | .u8LessThan τ e a b => .u8LessThan τ e (typedToSimple a) (typedToSimple b)
   | .u32LessThan τ e a b => .u32LessThan τ e (typedToSimple a) (typedToSimple b)
+  | .u8RangeCheck τ e a b => .u8RangeCheck τ e (typedToSimple a) (typedToSimple b)
+  | .toField τ e a => .toField τ e (typedToSimple a)
+  | .u8FromFieldUnsafe τ e a => .u8FromFieldUnsafe τ e (typedToSimple a)
   | .debug τ e l t r =>
     let t' := match t with | none => none | some sub => some (typedToSimple sub)
     .debug τ e l t' (typedToSimple r)

Ix/Aiur/Compiler/Simple.lean

@@ -108,6 +108,16 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
       let a' ← simplifyTypedTerm decls a
       let b' ← simplifyTypedTerm decls b
       pure (.u32LessThan τ e a' b')
+  | .u8RangeCheck τ e a b => do
+      let a' ← simplifyTypedTerm decls a
+      let b' ← simplifyTypedTerm decls b
+      pure (.u8RangeCheck τ e a' b')
+  | .toField τ e a => do
+      let a' ← simplifyTypedTerm decls a
+      pure (.toField τ e a')
+  | .u8FromFieldUnsafe τ e a => do
+      let a' ← simplifyTypedTerm decls a
+      pure (.u8FromFieldUnsafe τ e a')
   | t => pure t
 termination_by t => sizeOf t
 decreasing_by

Ix/Aiur/Interpret.lean

@@ -400,6 +400,15 @@ partial def interp (decls : Decls) (bindings : Bindings) : Term → InterpM Valu
       | .field a, .field b =>
           return .field (if a.val.toUInt32 < b.val.toUInt32 then 1 else 0)
       | _, _ => throwErr "u32LessThan: expected field values"
+  | .u8RangeCheck t1 t2 => do
+      match (← interp decls bindings t1), (← interp decls bindings t2) with
+      | .field a, .field b =>
+          if a.val < 256 && b.val < 256 then return .tuple #[.field a, .field b]
+          else throwErr "u8RangeCheck: value out of range [0, 256)"
+      | _, _ => throwErr "u8RangeCheck: expected field values"
+  -- `toField` / `u8FromFieldUnsafe` are erased coercions: value unchanged.
+  | .toField t | .u8FromFieldUnsafe t => interp decls bindings t
+  | .u8Lit n => return .field (G.ofNat n)
   | .debug label optT cont => do
       match optT with
       | none   => dbg_trace s!"{label}"