Aiur refactor (#289)

* Simplified types

No more `ContextualType`. Instead, `TypedInner` carries an `escape` field

* Infer term refactor

into many smaller functions

* And and Or U8 operations

* fmt and clippy

* bincode unmaintained

Author: gabriel-barrett

Ix/Aiur/Bytecode.lean

@@ -27,6 +27,8 @@ inductive Op
   | u8ShiftRight : ValIdx → Op
   | u8Xor : ValIdx → ValIdx → Op
   | u8Add : ValIdx → ValIdx → Op
+  | u8And : ValIdx → ValIdx → Op
+  | u8Or : ValIdx → ValIdx → Op
   | debug : String → Option (Array ValIdx) → Op
   deriving Repr
 

Ix/Aiur/Check.lean

@@ -120,7 +120,7 @@ def opLayout : Bytecode.Op → LayoutM Unit
     pushDegrees $ .replicate 8 1
     bumpAuxiliaries 8
     bumpLookups
-  | .u8ShiftLeft _ | .u8ShiftRight _ | .u8Xor .. => do
+  | .u8ShiftLeft _ | .u8ShiftRight _ | .u8Xor .. | .u8And .. | .u8Or .. => do
     pushDegree 1
     bumpAuxiliaries 1
     bumpLookups
@@ -246,7 +246,7 @@ partial def toIndex
  (term : TypedTerm) : StateM CompilerState (Array Bytecode.ValIdx) :=
   match term.inner with
   -- | .unsafeCast inner castTyp =>
-  --   if typSize layoutMap castTyp != typSize layoutMap term.typ.unwrap then
+  --   if typSize layoutMap castTyp != typSize layoutMap term.typ then
   --     panic! "Impossible cast"
   --   else
   --     toIndex layoutMap bindings (.mk term.typ inner)
@@ -335,32 +335,32 @@ partial def toIndex
   --     pushOp (Bytecode.Op.preimg layout.index out layout.inputSize) layout.inputSize
   --   | _ => panic! "should not happen after typechecking"
   | .proj arg i => do
-    let typs := match arg.typ with
-      | .evaluates (.tuple typs) => typs
+    let typs := match (arg.typ, arg.escapes) with
+      | (.tuple typs, false) => typs
       | _ => panic! "Should not happen after typechecking"
     let offset := (typs.extract 0 i).foldl (init := 0)
       fun acc typ => typSize layoutMap typ + acc
     let arg ← toIndex layoutMap bindings arg
     let length := typSize layoutMap typs[i]!
     pure $ arg.extract offset (offset + length)
   | .get arr i => do
-    let eltTyp := match arr.typ with
-      | .evaluates (.array typ _) => typ
+    let eltTyp := match (arr.typ, arr.escapes) with
+      | (.array typ _, false) => typ
       | _ => panic! "Should not happen after typechecking"
     let eltSize := typSize layoutMap eltTyp
     let offset := i * eltSize
     let arr ← toIndex layoutMap bindings arr
     pure $ arr.extract offset (offset + eltSize)
   | .slice arr i j => do
-    let eltTyp := match arr.typ with
-      | .evaluates (.array typ _) => typ
+    let eltTyp := match (arr.typ, arr.escapes) with
+      | (.array typ _, false) => typ
       | _ => panic! "Should not happen after typechecking"
     let eltSize := typSize layoutMap eltTyp
     let arr ← toIndex layoutMap bindings arr
     pure $ arr.extract (i * eltSize) (j * eltSize)
   | .set arr i val => do
-    let eltTyp := match arr.typ with
-      | .evaluates (.array typ _) => typ
+    let eltTyp := match (arr.typ, arr.escapes) with
+      | (.array typ _, false) => typ
       | _ => panic! "Should not happen after typechecking"
     let eltSize := typSize layoutMap eltTyp
     let arr ← toIndex layoutMap bindings arr
@@ -372,8 +372,8 @@ partial def toIndex
     let args ← toIndex layoutMap bindings arg
     pushOp (.store args)
   | .load ptr => do
-    let size := match ptr.typ.unwrap with
-    | .pointer typ => typSize layoutMap typ
+    let size := match (ptr.typ, ptr.escapes) with
+    | (.pointer typ, false) => typSize layoutMap typ
     | _ => unreachable!
     let ptr ← expectIdx ptr
     pushOp (.load size ptr) size
@@ -421,6 +421,14 @@ partial def toIndex
     let i ← expectIdx i
     let j ← expectIdx j
     pushOp (.u8Add i j) 2
+  | .u8And i j => do
+    let i ← expectIdx i
+    let j ← expectIdx j
+    pushOp (.u8And i j)
+  | .u8Or i j => do
+    let i ← expectIdx i
+    let j ← expectIdx j
+    pushOp (.u8Or i j)
   | .debug label term ret => do
     let term ← term.mapM (toIndex layoutMap bindings)
     modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
@@ -454,7 +462,7 @@ partial def TypedTerm.compile
     modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
     ret.compile returnTyp layoutMap bindings
   | .match term cases =>
-    match term.typ.unwrapOr returnTyp with
+    match term.typ with
     -- Also do this for tuple-like and array-like (one constructor only) datatypes
     | .tuple typs => match cases with
       | [(.tuple vars, branch)] => do

Ix/Aiur/Compile.lean

@@ -123,6 +123,8 @@ syntax "u8_shift_left" "(" trm ")"                            : trm
 syntax "u8_shift_right" "(" trm ")"                           : trm
 syntax "u8_xor" "(" trm ", " trm ")"                          : trm
 syntax "u8_add" "(" trm ", " trm ")"                          : trm
+syntax "u8_and" "(" trm ", " trm ")"                          : trm
+syntax "u8_or" "(" trm ", " trm ")"                           : trm
 syntax "dbg!" "(" str (", " trm)? ")" ";" (trm)?              : trm
 
 syntax trm "[" "@" noWs ident "]"                                                        : trm
@@ -219,6 +221,10 @@ partial def elabTrm : ElabStxCat `trm
     mkAppM ``Term.u8Xor #[← elabTrm i, ← elabTrm j]
   | `(trm| u8_add($i:trm, $j:trm)) => do
     mkAppM ``Term.u8Add #[← elabTrm i, ← elabTrm j]
+  | `(trm| u8_and($i:trm, $j:trm)) => do
+    mkAppM ``Term.u8And #[← elabTrm i, ← elabTrm j]
+  | `(trm| u8_or($i:trm, $j:trm)) => do
+    mkAppM ``Term.u8Or #[← elabTrm i, ← elabTrm j]
   | `(trm| dbg!($label:str $[, $t:trm]?); $[$ret:trm]?) => do
     let t ← match t with
       | none => mkAppOptM ``Option.none #[some (mkConst ``Term)]
@@ -360,6 +366,14 @@ where
       let i ← replaceToken old new i
       let j ← replaceToken old new j
       `(trm| u8_add($i, $j))
+    | `(trm| u8_and($i:trm, $j:trm)) => do
+      let i ← replaceToken old new i
+      let j ← replaceToken old new j
+      `(trm| u8_and($i, $j))
+    | `(trm| u8_or($i:trm, $j:trm)) => do
+      let i ← replaceToken old new i
+      let j ← replaceToken old new j
+      `(trm| u8_or($i, $j))
     | `(trm| dbg!($label:str $[, $t:trm]?); $[$ret:trm]?) => do
       let t' ← t.mapM $ replaceToken old new
       let ret' ← ret.mapM $ replaceToken old new

Ix/Aiur/Meta.lean

@@ -93,6 +93,8 @@ inductive Term
   | u8ShiftRight : Term → Term
   | u8Xor : Term → Term → Term
   | u8Add : Term → Term → Term
+  | u8And : Term → Term → Term
+  | u8Or : Term → Term → Term
   | debug : String → Option Term → Term → Term
   deriving Repr, BEq, Hashable, Inhabited
 
@@ -104,19 +106,6 @@ inductive Data
 
 end
 
-inductive ContextualType
-  | evaluates : Typ → ContextualType
-  | escapes : ContextualType
-  deriving Repr, BEq, Inhabited
-
-def ContextualType.unwrap : ContextualType → Typ
-| .escapes => panic! "term should not escape"
-| .evaluates typ => typ
-
-def ContextualType.unwrapOr : ContextualType → Typ → Typ
-| .escapes => fun typ => typ
-| .evaluates typ => fun _ => typ
-
 mutual
 inductive TypedTermInner
   | unit
@@ -149,12 +138,15 @@ inductive TypedTermInner
   | u8ShiftRight : TypedTerm → TypedTermInner
   | u8Xor : TypedTerm → TypedTerm → TypedTermInner
   | u8Add : TypedTerm → TypedTerm → TypedTermInner
+  | u8And : TypedTerm → TypedTerm → TypedTermInner
+  | u8Or : TypedTerm → TypedTerm → TypedTermInner
   | debug : String → Option TypedTerm → TypedTerm → TypedTermInner
   deriving Repr, Inhabited
 
 structure TypedTerm where
-  typ : ContextualType
+  typ : Typ
   inner : TypedTermInner
+  escapes : Bool
   deriving Repr, Inhabited
 
 inductive TypedData

Ix/Aiur/Term.lean

@@ -199,6 +199,14 @@ def toplevel := ⟦
     (u8_add(i_xor_j, i), u8_add(i_xor_j, j))
   }
 
+  fn u8_add_and(i: G, j: G) -> G {
+    u8_and(i, j)
+  }
+
+  fn u8_add_or(i: G, j: G) -> G {
+    u8_or(i, j)
+  }
+
   fn fold_matrix_sum(m: [[G; 2]; 2]) -> G {
     fold(0 .. 2, 0, |acc_outer, @i|
       fold(0 .. 2, acc_outer, |acc_inner, @j|
@@ -245,6 +253,8 @@ def aiurTestCases : List AiurTestCase := [
       ⟨#[1, 2, 3, 4, 1, 2, 3, 4], .ofList [(#[0], ⟨0, 4⟩), (#[1], ⟨0, 8⟩)]⟩⟩,
     .noIO `shr_shr_shl_decompose #[87] #[0, 1, 0, 1, 0, 1, 0, 0],
     .noIO `u8_add_xor #[45, 131] #[219, 0, 49, 1],
+    .noIO `u8_add_and #[45, 131] #[1],
+    .noIO `u8_add_or #[45, 131] #[175],
     .noIO `fold_matrix_sum #[1, 2, 3, 4] #[10],
   ]
 

Tests/Aiur.lean

@@ -74,6 +74,8 @@ ignore = [
     "RUSTSEC-2024-0384", # `instant` crate is unmaintained
     "RUSTSEC-2024-0370", # `proc-macro-error` crate is unmaintained
     "RUSTSEC-2023-0089", # `atomic-polyfill` crate is unmaintained
+    "RUSTSEC-2025-0141", # `bincode` crate is unmaintained
+
     #{ id = "RUSTSEC-0000-0000", reason = "you can specify a reason the advisory is ignored" },
     #"a-crate-that-is-yanked@0.1.1", # you can also ignore yanked crate versions if you wish
     #{ crate = "a-crate-that-is-yanked@0.1.1", reason = "you can specify why you are ignoring the yanked crate" },

deny.toml

@@ -53,6 +53,8 @@ pub enum Op {
   U8ShiftRight(ValIdx),
   U8Xor(ValIdx, ValIdx),
   U8Add(ValIdx, ValIdx),
+  U8And(ValIdx, ValIdx),
+  U8Or(ValIdx, ValIdx),
   Debug(String, Option<Vec<ValIdx>>),
 }
 

src/aiur/bytecode.rs

@@ -16,8 +16,8 @@ use crate::aiur::{
     bytes1::{Bytes1, Bytes1Op},
     bytes2::{Bytes2, Bytes2Op},
   },
-  memory_channel, u8_add_channel, u8_bit_decomposition_channel,
-  u8_shift_left_channel, u8_shift_right_channel, u8_xor_channel,
+  memory_channel, u8_add_channel, u8_and_channel, u8_bit_decomposition_channel,
+  u8_or_channel, u8_shift_left_channel, u8_shift_right_channel, u8_xor_channel,
 };
 
 type Expr = SymbolicExpression<G>;
@@ -415,6 +415,22 @@ impl Op {
         sel.clone(),
         state,
       ),
+      Op::U8And(i, j) => bytes2_constraints(
+        *i,
+        *j,
+        &Bytes2Op::And,
+        u8_and_channel(),
+        sel.clone(),
+        state,
+      ),
+      Op::U8Or(i, j) => bytes2_constraints(
+        *i,
+        *j,
+        &Bytes2Op::Or,
+        u8_or_channel(),
+        sel.clone(),
+        state,
+      ),
       Op::IOSetInfo(..) | Op::IOWrite(_) | Op::Debug(..) => (),
     }
   }

src/aiur/constraints.rs

@@ -271,6 +271,20 @@ impl Function {
             bytes2_execute(*i, *j, &Bytes2Op::Add, &mut map, record);
           }
         },
+        ExecEntry::Op(Op::U8And(i, j)) => {
+          if unconstrained {
+            map.push(Bytes2::and(&map[*i], &map[*j]));
+          } else {
+            bytes2_execute(*i, *j, &Bytes2Op::And, &mut map, record);
+          }
+        },
+        ExecEntry::Op(Op::U8Or(i, j)) => {
+          if unconstrained {
+            map.push(Bytes2::or(&map[*i], &map[*j]));
+          } else {
+            bytes2_execute(*i, *j, &Bytes2Op::Or, &mut map, record);
+          }
+        },
         ExecEntry::Op(Op::Debug(label, idxs)) => match idxs {
           None => println!("{label}"),
           Some(idxs) => {