Merge pull request #8 from lindy-labs/selfref

Add support for self referential types

Aegis/Analyzer.lean

@@ -7,67 +7,6 @@ open Lean Expr Meta Qq
 
 namespace Sierra
 
-set_option trace.aesop.extraction true
-
-def buildTypeDefs (typedefs : List (Identifier × Identifier)) :
-    Except String (HashMap Identifier SierraType) := do
-  let mut acc := HashMap.empty
-  for (name, ty) in typedefs do
-    let v : SierraType ← go acc ty
-    acc := acc.insert name v
-  return acc
-where go (acc : _) (ty : Identifier) : Except String SierraType :=
-  match ty with
-  | .name "felt252" [] .none => pure .Felt252
-  | .name "u8" [] .none => pure .U8
-  | .name "u16" [] .none => pure .U16
-  | .name "u32" [] .none => pure .U32
-  | .name "u64" [] .none => pure .U64
-  | .name "u128" [] .none => pure .U128
-  | .name "RangeCheck" [] .none => pure .RangeCheck
-  | .name "Enum" (.usertype _ :: l) .none => do
-    let l ← flip mapM l fun x => match x with
-      | .identifier ident => pure ident
-      | _ => throw "Expected Enum parameter to refer a to a type"
-    pure <| .Enum (l.map acc.find!)
-  | .name "Struct" (.usertype _ :: l) .none => do
-    let l ← flip mapM l fun x => match x with
-      | .identifier ident => pure ident
-      | _ => throw "Expected Enum parameter to refer a to a type"
-    pure <| .Struct (l.map acc.find!)
-  | .name "NonZero" (Parameter.identifier ident :: []) .none => do
-    pure <| .NonZero <| acc.find! ident
-  | .name "Box" [l] .none =>
-    match l with
-    | .identifier ident => pure <| .Box <| acc.find! ident
-    | _ => throw "Expected Box parameter to refer to a type"
-  | .name "Snapshot" [l] .none =>
-    match l with
-    | .identifier ident => pure <| .Snapshot <| acc.find! ident
-    | _ => throw "Expected Snapshot parameter to refer to a type"
-  | .name "Array" [t] .none =>
-    match t with
-    | .identifier ident => pure <| .Array <| acc.find! ident
-    | _ => throw "Expected ARray parameter to refer to a type"
-  | .name "U128MulGuarantee" [] .none => pure .U128MulGuarantee
-  | .name "Pedersen" [] .none => pure .Pedersen
-  | .name "BuiltinCosts" [] .none => pure .BuiltinCosts
-  | .name "GasBuiltin" [] .none => pure .GasBuiltin
-  | .name "Bitwise" [] .none => pure .Bitwise
-  | .name "Uninitialized" [t] .none =>
-    match t with
-    | .identifier ident => pure <| .Array <| acc.find! ident
-    | _ => throw "Expected Uninitalized parameter to refer to a type"
-  | .name "Nullable" [t] .none =>
-    match t with
-    | .identifier ident => pure <| .Nullable <| acc.find! ident
-    | _ => throw "Expected Nullable parameter to refer to a type"
-  | .name "StorageBaseAddress" [] .none => pure .StorageBaseAddress
-  | .name "StorageAddress" [] .none => pure .StorageAddress
-  | .name "System" [] .none => pure .System
-  | .name "ContractAddress" [] .none => pure .ContractAddress
-  | _ => throw s!"Unhandled type {ty}"
-
 def buildFuncSignatures
   (currentFunc : Identifier)
   (typedefs : HashMap Identifier SierraType)
@@ -79,7 +18,7 @@ def buildFuncSignatures
   for (name, sig) in funcdefs do
     match FuncData.libfuncs currentFunc typedefs specs metadataRef sig with
     | some sig => acc := acc.insert name sig
-    | none => () --throw <| toString name ++ " : no such libfunc"
+    | none => dbg_trace s!"{toString name} : no such libfunc"
   return acc
 
 structure State where
@@ -164,7 +103,7 @@ partial def processState
     let s ← get
     let es := (s.outputRefs.zip (st.args.map fun n => s.refs.find! n)).zip s.outputTypes
     let es := es.map fun ((ofv, rfv), t) =>
-      Expr.mkEq (SierraType.toQuote <| typeDefs.find! t) (.fvar rfv) (.fvar ofv)
+      Expr.mkEq (SierraType.toQuote [] <| typeDefs.find! t) (.fvar rfv) (.fvar ofv)
     return (st, .cons (Expr.mkAnds es) [])
   | _ => do
     let .some st := f.statements.get? (← get).pc
@@ -231,7 +170,7 @@ partial def getFuncCondition (ident : Identifier) (pc : ℕ) (inputArgs : List (
   for t in outputTypes do
     let name ← mkFreshUserName `ρ
     let fv ← mkFreshFVarId
-    lctx := lctx.mkLocalDecl fv name <| SierraType.toQuote <| typeDefs.find! t
+    lctx := lctx.mkLocalDecl fv name <| SierraType.toQuote [] <| typeDefs.find! t
     outputRefs := outputRefs.push fv
   -- Create fvar for the reference to the `Metadata` instance
   let metadataRef ← mkFreshFVarId
@@ -249,7 +188,7 @@ partial def getFuncCondition (ident : Identifier) (pc : ℕ) (inputArgs : List (
     let mut refs : RefTable := ∅
     -- Add input arguments to initial local context and refs table
     for (i, t) in inputArgs do
-      refs := refs.insert i <| ← getOrMkNewRef i <| SierraType.toQuote <| typeDefs.find! t
+      refs := refs.insert i <| ← getOrMkNewRef i <| SierraType.toQuote [] <| typeDefs.find! t
     set { (← get) with refs := refs }
     let (_, cs) ← processState typeDefs funcSigs sf inputArgs
     processAndOrTree inputArgs cs) s
@@ -279,10 +218,10 @@ def getLocalDeclInfos (ident : Identifier) (pc : ℕ) (inputArgs : List (ℕ ×
   ret := ret.push <| .mk (← mkFreshUserName `m) fun _ => pure q(Metadata)
   for (_, t) in inputArgs do
     let n ← mkFreshUserName `a  -- TODO make anonymous?
-    ret := ret.push <| .mk n fun _ => pure <| SierraType.toQuote <| typeDefs.find! t
+    ret := ret.push <| .mk n fun _ => pure <| SierraType.toQuote [] <| typeDefs.find! t
   for t in outputTypes do
     let n ← mkFreshUserName `ρ  -- TODO make anonymous?
-    ret := ret.push <| .mk n fun _ => pure <| SierraType.toQuote <| typeDefs.find! t
+    ret := ret.push <| .mk n fun _ => pure <| SierraType.toQuote [] <| typeDefs.find! t
   let n ← mkFreshUserName `h_auto
   let e ← getFuncCondition sf specs ident pc inputArgs outputTypes
   -- Add the auto spec, to which all previous args are applied

Aegis/Libfuncs.lean

@@ -28,7 +28,7 @@ namespace Sierra.FuncData
 
 variable (currentFunc : Identifier) (typeRefs : HashMap Identifier SierraType)
   (specs : HashMap Identifier (Name × (FVarId → FuncData)))  (metadataRef : FVarId)
-  (i : Identifier) 
+  (i : Identifier)
 
 /-- The definition of `libfuncs` is split into pieces do to slow elaboration time. -/
 private def libfuncs_aux :=
@@ -38,13 +38,13 @@ private def libfuncs_aux :=
   <|> boolLibfuncs i
   <|> enumLibfuncs typeRefs i
   <|> structLibfuncs typeRefs i
-  <|> boxLibfuncs typeRefs i
+  <|> boxLibfuncs metadataRef typeRefs i
   <|> snapshotLibfuncs typeRefs i
   <|> nonZeroLibfuncs typeRefs i
 
 private def libfuncs_aux2 :=
-  libfuncs_aux typeRefs i
-  <|> arrayLibfuncs typeRefs i
+  libfuncs_aux typeRefs metadataRef i
+  <|> arrayLibfuncs metadataRef typeRefs i
   <|> functionCallLibfuncs specs metadataRef i
   <|> syscallLibfuncs metadataRef i
   <|> builtinCostsLibfuncs currentFunc metadataRef i

Aegis/Libfuncs/Array.lean

@@ -4,6 +4,8 @@ open Qq Lean Sierra
 
 namespace Sierra.FuncData
 
+variable (metadataRef : FVarId)
+
 def array_new (t : SierraType) : FuncData where
   inputTypes := []
   branches := [{ outputTypes := [.Array t]
@@ -12,17 +14,26 @@ def array_new (t : SierraType) : FuncData where
 def array_append (t : SierraType) : FuncData where
   inputTypes := [.Array t, t]
   branches := [{ outputTypes := [.Array t]
-                 condition := fun (a : Q(List $t.toQuote)) (b : Q($t.toQuote)) 
+                 condition := fun (a : Q(List $t.toQuote)) (b : Q($t.toQuote))
                    (ρ : Q(List $t.toQuote)) => q($ρ = $a ++ [$b]) }]
 
+example (f : Bool → Q(Type)) (b : Bool) (a b : Q(List $(f b))) : Q(Prop) :=
+  q(∃ x, x :: $a = $b)
+
+def array_pop_front_aux (T : Q(Type)) (m' : Q(Option $T)) (a ρ₁ : Q(List $T)) : Q(Prop) :=
+  q(∃ ρ₂', $m' = .some ρ₂' ∧ ρ₂' :: $ρ₁ = $a)
+
 def array_pop_front (t : SierraType) : FuncData where
   inputTypes := [.Array t]
   branches := [{ outputTypes := [.Array t, .Box t]
-                 condition := fun (a ρ₁: Q(List $t.toQuote)) (ρ₂ : Q($t.toQuote)) =>
-                   q($ρ₂ :: $ρ₁ = $a) },
+                 condition :=
+                   fun a ρ₁ (ρ₂ : Q(Nat)) =>
+                   let m : Q(Metadata) := .fvar metadataRef
+                   let m' : Expr := q($(m).boxHeap $t $ρ₂)
+                   array_pop_front_aux ⟦t⟧ m' a ρ₁ },
                { outputTypes := [.Array t]
                  condition := fun (a : Q(List $t.toQuote)) (ρ : Q(List $t.toQuote)) =>
-                   q($a = [] ∧ $ρ = []) }] -- TODO not actually clear what `ρ` is
+                   q($a = [] ∧ $ρ = []) }] -- TODO maybe add possibility of box being empty
 
 def array_len (t : SierraType) : FuncData where
   inputTypes := [.Snapshot (.Array t)]
@@ -54,7 +65,7 @@ def arrayLibfuncs (typeRefs : HashMap Identifier SierraType) : Identifier → Op
 | .name "array_append" [.identifier ident] .none =>
   return array_append (← typeRefs.find? ident)
 | .name "array_pop_front" [.identifier ident] .none =>
-  return array_pop_front (← typeRefs.find? ident)
+  return array_pop_front metadataRef (← typeRefs.find? ident)
 | .name "array_len" [.identifier ident] .none =>
   return array_len (← typeRefs.find? ident)
 | .name "array_get" [.identifier ident] .none =>

Aegis/Libfuncs/Box.lean

@@ -4,21 +4,33 @@ open Qq Lean Sierra
 
 namespace Sierra.FuncData
 
-def box_into (t : SierraType) : FuncData where
+variable (metadataRef : FVarId)
+
+def into_box (t : SierraType) : FuncData where
   inputTypes := [t]
-  branches := [{ outputTypes := [.Box t], condition := fun (a ρ : Q($(⟦t⟧))) => q($ρ = $a) }]
+  branches := [{ outputTypes := [.Box t]
+                 condition := fun a (ρ : Q(Nat)) =>
+                   let m : Q(Metadata) := .fvar metadataRef
+                   let lhs : Q(Option ($(t).toType)) := q($(m).boxHeap $t $ρ)
+                   let rhs := mkAppN (mkConst `Option.some [levelZero]) #[t.toQuote, a]
+                   Expr.mkEq q(Option ($(t).toType)) lhs rhs }]
 
 def unbox (t : SierraType) : FuncData where
   inputTypes := [.Box t]
-  branches := [{ outputTypes := [t], condition := fun (a ρ : Q($(⟦t⟧))) => q($ρ = $a) }]
+  branches := [{ outputTypes := [t]
+                 condition := fun (a : Q(Nat)) (ρ : Expr) =>
+                   let m : Q(Metadata) := .fvar metadataRef
+                   let rhs : Q(Option ($(t).toType)) := q($(m).boxHeap $t $a)
+                   let lhs := mkAppN (mkConst `Option.some [levelZero]) #[t.toQuote, ρ]
+                   Expr.mkEq q(Option ($(t).toType)) lhs rhs }]
 
 def boxLibfuncs (typeRefs : HashMap Identifier SierraType) : Identifier → Option FuncData
-| .name "box_into" [.identifier ident] _ =>
-  match typeRefs.find? ident with
-  | .some t => box_into t
+| .name "into_box" [.identifier ident] _ =>
+  match getMuBody <$> typeRefs.find? ident with
+  | .some t => into_box metadataRef t
   | _ => .none
 | .name "unbox" [.identifier ident] _ =>
-  match typeRefs.find? ident with
-  | .some t => unbox t
+  match getMuBody <$> typeRefs.find? ident with
+  | .some t => unbox metadataRef t
   | _ => .none
 | _ => .none

Aegis/Libfuncs/Enum.lean

@@ -39,17 +39,19 @@ def enum_snapshot_match (fields : List SierraType) : FuncData where
 
 def enumLibfuncs (typeRefs : HashMap Identifier SierraType) : Identifier → Option FuncData
 | .name "enum_init" [.identifier ident, .const (.ofNat n)] .none =>
-  match typeRefs.find? ident with
+  match getMuBody <$> typeRefs.find? ident with
   | .some (.Enum fields) =>
     if hn : n < fields.length then enum_init fields ⟨n, hn⟩
     else .none
   | _ => .none
 | .name "enum_match" [.identifier ident] .none =>
-  match typeRefs.find? ident with
-  | .some (.Enum fields) => enum_match fields
+  match getMuBody <$> typeRefs.find? ident with
+  | .some (.Enum fields) =>
+    enum_match fields
   | _ => .none
 | .name "enum_snapshot_match" [.identifier ident] .none =>
-  match typeRefs.find? ident with
-  | .some (.Enum fields) => enum_snapshot_match fields
+  match getMuBody <$> typeRefs.find? ident with
+  | .some (.Enum fields) =>
+    enum_snapshot_match fields
   | _ => .none
 | _ => .none

Aegis/Libfuncs/Struct.lean

@@ -47,15 +47,15 @@ def struct_snapshot_deconstruct (fields : List SierraType) : FuncData where
 
 def structLibfuncs (typeRefs : HashMap Identifier SierraType) : Identifier → Option FuncData
 | .name "struct_construct" [.identifier ident] _ =>
-  match typeRefs.find? ident with
+  match getMuBody <$> typeRefs.find? ident with
   | .some (.Struct fields) => struct_construct fields
   | _ => .none
 | .name "struct_deconstruct" [.identifier ident] _ =>
-  match typeRefs.find? ident with
+  match getMuBody <$> typeRefs.find? ident with
   | .some (.Struct fields) => struct_deconstruct fields
   | _ => .none
 | .name "struct_snapshot_deconstruct" [.identifier ident] _ =>
-  match typeRefs.find? ident with
+  match getMuBody <$> typeRefs.find? ident with
   | .some (.Struct fields) => struct_snapshot_deconstruct fields
   | _ => .none
 | _ => .none

Aegis/Libfuncs/UInt128.lean

@@ -9,7 +9,7 @@ namespace Sierra.FuncData
 def u128_overflowing_add : FuncData where
   inputTypes := [RangeCheck, U128, U128]
   branches := [{ outputTypes := [RangeCheck, U128]
-                 condition := fun _ (a b : Q(UInt128)) _ (ρ : Q(UInt128)) => 
+                 condition := fun _ (a b : Q(UInt128)) _ (ρ : Q(UInt128)) =>
                    q(($a).val + ($b).val < U128_MOD ∧ $ρ = $a + $b) },
                -- TODO check branch order
                { outputTypes := [RangeCheck, U128]
@@ -19,7 +19,7 @@ def u128_overflowing_add : FuncData where
 def u128_overflowing_sub : FuncData where
   inputTypes := [RangeCheck, U128, U128]
   branches := [{ outputTypes := [RangeCheck, U128]
-                 condition := fun _ (a b : Q(UInt128)) _ (ρ : Q(UInt128)) => 
+                 condition := fun _ (a b : Q(UInt128)) _ (ρ : Q(UInt128)) =>
                    q(($b).val ≤ ($a).val ∧ $ρ = $a - $b) },
                -- TODO check branch order
                { outputTypes := [RangeCheck, U128]
@@ -81,9 +81,9 @@ def bitwise : FuncData where
   inputTypes := [Bitwise, U128, U128]
   branches := [{ outputTypes := [Bitwise, U128, U128, U128]
                  condition := fun _ (lhs rhs : Q(UInt128)) _ (and xor or : Q(UInt128)) =>
-                   q($(and).val = Nat.land' $(lhs).val $(rhs).val
-                     ∧ $(xor).val = Nat.lxor' $(lhs).val $(rhs).val
-                     ∧ $(or).val = Nat.lor' $(lhs).val $(rhs).val) }]
+                   q($(and).val = Nat.land $(lhs).val $(rhs).val
+                     ∧ $(xor).val = Nat.xor $(lhs).val $(rhs).val
+                     ∧ $(or).val = Nat.lor $(lhs).val $(rhs).val) }]
 
 def uint128Libfuncs : Identifier → Option FuncData
 | .name "u128_overflowing_add" [] .none      => u128_overflowing_add

Aegis/Tests/ArrayPopFront.lean

@@ -36,10 +36,15 @@ test::pop_front@0([0]: Array<felt252>) -> (Array<felt252>, core::option::Option:
 
 aegis_info "test::pop_front"
 
-aegis_spec "test::pop_front" := fun _ a ρ₁ ρ₂ =>
-  ρ₁ = a.tail ∧ ρ₂ = if a.length = 0 then .inr () else .inl a.head!
-
-aegis_prove "test::pop_front" := fun _ a ρ₁ ρ₂ => by
-  rintro ⟨_,_,(⟨rfl,rfl,rfl⟩|⟨rfl,rfl,rfl⟩)⟩ <;> simp [«spec_test::pop_front»]
+aegis_spec "test::pop_front" := fun m a ρ₁ ρ₂ =>
+  ρ₁ = a.tail
+  ∧ (a.length = 0 ∧ ρ₂ = .inr ()
+    ∨ ∃ ρ₂', ρ₂ = .inl ρ₂' ∧ m.boxHeap .Felt252 ρ₂' = .some a.head!)
+
+aegis_prove "test::pop_front" := fun m a ρ₁ ρ₂ => by
+  unfold «spec_test::pop_front»
+  rintro ⟨_,_,(⟨⟨ρ₂',h,rfl⟩,rfl,rfl⟩|⟨rfl,rfl,rfl⟩)⟩
+  · aesop
+  · simp
 
 aegis_complete

Aegis/Tests/Commands.lean

@@ -128,6 +128,43 @@ aegis_spec "test::u128_const" := fun _ _ => True
 
 aegis_prove "test::u128_const" := fun _ _ _ => True.intro
 
+aegis_load_string "type felt252 = felt252 [storable: true, drop: true, dup: true, zero_sized: false];
+type Box<felt252> = Box<felt252> [storable: true, drop: true, dup: true, zero_sized: false];
+
+libfunc into_box<felt252> = into_box<felt252>;
+
+into_box<felt252>([0]) -> ([1]); // 0
+return([1]); // 1
+
+test::into_box@0([0]: felt252) -> (Box<felt252>);"
+
+aegis_spec "test::into_box" := fun m a ρ =>
+  m.boxHeap .Felt252 ρ = .some a
+
+aegis_prove "test::into_box" := fun _ _ _ _ => by
+  aesop
+
+aegis_load_string "type Box<felt252> = Box<felt252> [storable: true, drop: true, dup: true, zero_sized: false];
+type felt252 = felt252 [storable: true, drop: true, dup: true, zero_sized: false];
+
+libfunc unbox<felt252> = unbox<felt252>;
+libfunc store_temp<felt252> = store_temp<felt252>;
+
+unbox<felt252>([0]) -> ([1]); // 0
+store_temp<felt252>([1]) -> ([1]); // 1
+return([1]); // 2
+
+test::unbox@0([0]: Box<felt252>) -> (felt252);"
+
+aegis_spec "test::unbox" := fun m a ρ =>
+  m.boxHeap .Felt252 a = .some ρ
+
+aegis_prove "test::unbox" := fun m a b => by
+  unfold «spec_test::unbox»
+  rintro ⟨_,h,rfl⟩
+  symm
+  assumption
+
 aegis_load_string "type Unit = Struct<ut@Tuple>;
 type core::bool = Enum<ut@core::bool, Unit, Unit>;
 

Aegis/Tests/SelfRef.lean

@@ -0,0 +1,16 @@
+import Aegis.Commands
+import Aegis.Macros
+
+open Sierra Lean
+
+aegis_load_file "Aegis/Tests/issue2249.sierra"
+
+aegis_spec "issue2249::issue2249::simple_test" :=
+  fun _ ρ => ρ = .inl (12, .inr ())
+
+aegis_prove "issue2249::issue2249::simple_test" :=
+  fun _ ρ => by
+  unfold «spec_issue2249::issue2249::simple_test»
+  aesop
+
+aegis_complete