Transient storage fields (#1995)

Compiler/CompilationModel/Compile.lean

@@ -169,19 +169,20 @@ def compileStmt (fields : List Field) (events : List EventDef := [])
   | Stmt.setStorageWord field wordOffset value =>
       match findFieldWithResolvedSlot fields field with
       | some (f, slot) => do
+          let storeBuiltin := if f.isTransient then "tstore" else "sstore"
           let valueExpr ← compileExprWithInternals fields dynamicSource internalFunctions value
           let slotExpr (baseSlot : Nat) :=
             if wordOffset == 0 then YulExpr.lit baseSlot
             else YulExpr.call "add" [YulExpr.lit baseSlot, YulExpr.lit wordOffset]
           match f.aliasSlots with
           | [] =>
-              pure [YulStmt.expr (YulExpr.call "sstore" [slotExpr slot, valueExpr])]
+              pure [YulStmt.expr (YulExpr.call storeBuiltin [slotExpr slot, valueExpr])]
           | _ =>
               pure [
                 YulStmt.block (
                   [YulStmt.let_ "__compat_value" valueExpr] ++
                   (slot :: f.aliasSlots).map (fun writeSlot =>
-                    YulStmt.expr (YulExpr.call "sstore" [
+                    YulStmt.expr (YulExpr.call storeBuiltin [
                       slotExpr writeSlot,
                       YulExpr.ident "__compat_value"
                     ]))
@@ -200,19 +201,23 @@ def compileStmt (fields : List Field) (events : List EventDef := [])
         (← compileExprWithInternals fields dynamicSource internalFunctions key)
         (← compileExprWithInternals fields dynamicSource internalFunctions value)
         "setMapping"
+        0
+        true
   | Stmt.setMappingWord field key wordOffset value => do
       compileMappingSlotWrite fields field
         (← compileExprWithInternals fields dynamicSource internalFunctions key)
         (← compileExprWithInternals fields dynamicSource internalFunctions value)
         "setMappingWord"
         wordOffset
+        true
   | Stmt.setMappingPackedWord field key wordOffset packed value => do
       compileMappingPackedSlotWrite fields field
         (← compileExprWithInternals fields dynamicSource internalFunctions key)
         (← compileExprWithInternals fields dynamicSource internalFunctions value)
         wordOffset
         packed
         "setMappingPackedWord"
+        true
   | Stmt.setMapping2 field key1 key2 value =>
       compileSetMapping2 fields dynamicSource field key1 key2 value internalFunctions
   | Stmt.setMapping2Word field key1 key2 wordOffset value =>
@@ -222,6 +227,8 @@ def compileStmt (fields : List Field) (events : List EventDef := [])
         (← compileExprWithInternals fields dynamicSource internalFunctions key)
         (← compileExprWithInternals fields dynamicSource internalFunctions value)
         "setMappingUint"
+        0
+        true
   | Stmt.setMappingChain field keys value =>
       compileSetMappingChain fields dynamicSource field keys value internalFunctions
   | Stmt.setStructMember field key memberName value =>

Compiler/CompilationModel/ExpressionCompile.lean

@@ -28,11 +28,12 @@ def compileMappingSlotRead (fields : List Field) (field : String) (keyExpr : Yul
   if !isMapping fields field then
     throw s!"Compilation error: field '{field}' is not a mapping"
   else
-    match findFieldSlot fields field with
-    | some slot =>
+    match findFieldWithResolvedSlot fields field with
+    | some (f, slot) =>
       let mappingBase := YulExpr.call "mappingSlot" [YulExpr.lit slot, keyExpr]
       let finalSlot := if wordOffset == 0 then mappingBase else YulExpr.call "add" [mappingBase, YulExpr.lit wordOffset]
-      pure (YulExpr.call "sload" [finalSlot])
+      let loadBuiltin := if f.isTransient then "tload" else "sload"
+      pure (YulExpr.call loadBuiltin [finalSlot])
     | none => throw s!"Compilation error: unknown mapping field '{field}' in {label}"
 
 -- Exposed so proof modules can name the exact nested mapping-chain lowering shape.
@@ -191,12 +192,13 @@ def compileExprWithInternals (fields : List Field)
     else
       match findFieldWithResolvedSlot fields field with
       | some (f, slot) =>
+          let loadBuiltin := if f.isTransient then "tload" else "sload"
           match f.packedBits with
           | none =>
-              pure (YulExpr.call "sload" [YulExpr.lit slot])
+              pure (YulExpr.call loadBuiltin [YulExpr.lit slot])
           | some packed =>
               pure (YulExpr.call "and" [
-                YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call "sload" [YulExpr.lit slot]],
+                YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call loadBuiltin [YulExpr.lit slot]],
                 YulExpr.lit (packedMaskNat packed)
               ])
       | none => throw s!"Compilation error: unknown storage field '{field}'"
@@ -206,14 +208,15 @@ def compileExprWithInternals (fields : List Field)
     else
       match findFieldWithResolvedSlot fields field with
       | some (f, slot) =>
+          let loadBuiltin := if f.isTransient then "tload" else "sload"
           match f.ty with
           | .address =>
               match f.packedBits with
               | none =>
-                  pure (YulExpr.call "sload" [YulExpr.lit slot])
+                  pure (YulExpr.call loadBuiltin [YulExpr.lit slot])
               | some packed =>
                   pure (YulExpr.call "and" [
-                    YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call "sload" [YulExpr.lit slot]],
+                    YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call loadBuiltin [YulExpr.lit slot]],
                     YulExpr.lit (packedMaskNat packed)
                   ])
           | _ =>
@@ -236,36 +239,39 @@ def compileExprWithInternals (fields : List Field)
     if !isMapping2 fields field then
       throw s!"Compilation error: field '{field}' is not a double mapping"
     else
-      match findFieldSlot fields field with
-      | some slot => do
+      match findFieldWithResolvedSlot fields field with
+      | some (f, slot) => do
+        let loadBuiltin := if f.isTransient then "tload" else "sload"
         let key1Expr ← compileExprWithInternals fields dynamicSource internalFunctions key1
         let key2Expr ← compileExprWithInternals fields dynamicSource internalFunctions key2
         let innerSlot := YulExpr.call "mappingSlot" [YulExpr.lit slot, key1Expr]
-        pure (YulExpr.call "sload" [YulExpr.call "mappingSlot" [innerSlot, key2Expr]])
+        pure (YulExpr.call loadBuiltin [YulExpr.call "mappingSlot" [innerSlot, key2Expr]])
       | none => throw s!"Compilation error: unknown mapping field '{field}'"
   | Expr.mapping2Word field key1 key2 wordOffset =>
     if !isMapping2 fields field then
       throw s!"Compilation error: field '{field}' is not a double mapping"
     else
-      match findFieldSlot fields field with
-      | some slot => do
+      match findFieldWithResolvedSlot fields field with
+      | some (f, slot) => do
+        let loadBuiltin := if f.isTransient then "tload" else "sload"
         let key1Expr ← compileExprWithInternals fields dynamicSource internalFunctions key1
         let key2Expr ← compileExprWithInternals fields dynamicSource internalFunctions key2
         let innerSlot := YulExpr.call "mappingSlot" [YulExpr.lit slot, key1Expr]
         let outerSlot := YulExpr.call "mappingSlot" [innerSlot, key2Expr]
         let finalSlot := if wordOffset == 0 then outerSlot else YulExpr.call "add" [outerSlot, YulExpr.lit wordOffset]
-        pure (YulExpr.call "sload" [finalSlot])
+        pure (YulExpr.call loadBuiltin [finalSlot])
       | none => throw s!"Compilation error: unknown mapping field '{field}'"
   | Expr.mappingUint field key => do
       compileMappingSlotRead fields field (← compileExprWithInternals fields dynamicSource internalFunctions key) "mappingUint"
   | Expr.mappingChain field keys =>
       if !isMapping fields field then
         throw s!"Compilation error: field '{field}' is not a mapping"
       else
-        match findFieldSlot fields field with
-        | some slot => do
+        match findFieldWithResolvedSlot fields field with
+        | some (f, slot) => do
             let keyExprs ← compileExprListWithInternals fields dynamicSource internalFunctions keys
-            pure (YulExpr.call "sload" [compileMappingSlotChain (YulExpr.lit slot) keyExprs])
+            let loadBuiltin := if f.isTransient then "tload" else "sload"
+            pure (YulExpr.call loadBuiltin [compileMappingSlotChain (YulExpr.lit slot) keyExprs])
         | none => throw s!"Compilation error: unknown mapping field '{field}'"
   | Expr.structMember field key memberName => do
       if isMapping2 fields field then
@@ -295,19 +301,20 @@ def compileExprWithInternals (fields : List Field)
           match findStructMember members memberName with
           | none => throw s!"Compilation error: struct field '{field}' has no member '{memberName}'"
           | some member =>
-            match findFieldSlot fields field with
-            | some slot => do
+            match findFieldWithResolvedSlot fields field with
+            | some (f, slot) => do
+              let loadBuiltin := if f.isTransient then "tload" else "sload"
               let key1Expr ← compileExprWithInternals fields dynamicSource internalFunctions key1
               let key2Expr ← compileExprWithInternals fields dynamicSource internalFunctions key2
               let innerSlot := YulExpr.call "mappingSlot" [YulExpr.lit slot, key1Expr]
               let outerSlot := YulExpr.call "mappingSlot" [innerSlot, key2Expr]
               let finalSlot := if member.wordOffset == 0 then outerSlot else YulExpr.call "add" [outerSlot, YulExpr.lit member.wordOffset]
               match member.packed with
               | none =>
-                pure (YulExpr.call "sload" [finalSlot])
+                pure (YulExpr.call loadBuiltin [finalSlot])
               | some packed =>
                 pure (YulExpr.call "and" [
-                  YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call "sload" [finalSlot]],
+                  YulExpr.call "shr" [YulExpr.lit packed.offset, YulExpr.call loadBuiltin [finalSlot]],
                   YulExpr.lit (packedMaskNat packed)
                 ])
             | none => throw s!"Compilation error: unknown mapping field '{field}'"

Compiler/CompilationModel/MappingWrites.lean

@@ -7,48 +7,81 @@ open Compiler
 open Compiler.Yul
 
 def compileMappingSlotWrite (fields : List Field) (field : String)
-    (keyExpr valueExpr : YulExpr) (label : String) (wordOffset : Nat := 0) : Except String (List YulStmt) :=
+    (keyExpr valueExpr : YulExpr) (label : String) (wordOffset : Nat := 0)
+    (allowTransient : Bool := false) : Except String (List YulStmt) :=
   if !isMapping fields field then
     throw s!"Compilation error: field '{field}' is not a mapping"
   else
+    let storeBuiltin :=
+      if allowTransient then
+        match findFieldWithResolvedSlot fields field with
+        | some (f, _) => if f.isTransient then "tstore" else "sstore"
+        | none => "sstore"
+      else
+        "sstore"
     match findFieldWriteSlots fields field with
     | some slots =>
         match slots with
-        | [] =>
-            throw s!"Compilation error: internal invariant failure: no write slots for mapping field '{field}' in {label}"
-        | [slot] =>
-            let mappingBase := YulExpr.call "mappingSlot" [YulExpr.lit slot, keyExpr]
-            let writeSlot := if wordOffset == 0 then mappingBase else YulExpr.call "add" [mappingBase, YulExpr.lit wordOffset]
-            pure [
-              YulStmt.expr (YulExpr.call "sstore" [
-                writeSlot,
-                valueExpr
-              ])
-            ]
-        | _ =>
-            let compatSlotExpr := fun (slot : Nat) =>
-              let mappingBase := YulExpr.call "mappingSlot" [YulExpr.lit slot, YulExpr.ident "__compat_key"]
-              if wordOffset == 0 then mappingBase else YulExpr.call "add" [mappingBase, YulExpr.lit wordOffset]
-            pure [
-              YulStmt.block (
-                [YulStmt.let_ "__compat_key" keyExpr, YulStmt.let_ "__compat_value" valueExpr] ++
-                slots.map (fun slot =>
+          | [] =>
+              throw s!"Compilation error: internal invariant failure: no write slots for mapping field '{field}' in {label}"
+          | [slot] =>
+              let mappingBase := YulExpr.call "mappingSlot" [YulExpr.lit slot, keyExpr]
+              let writeSlot := if wordOffset == 0 then mappingBase else YulExpr.call "add" [mappingBase, YulExpr.lit wordOffset]
+              if allowTransient then
+                match findFieldWithResolvedSlot fields field with
+                | some (_, _) =>
+                    pure [
+                      YulStmt.expr (YulExpr.call storeBuiltin [
+                        writeSlot,
+                        valueExpr
+                      ])
+                    ]
+                | none => throw s!"Compilation error: unknown mapping field '{field}' in {label}"
+              else
+                pure [
                   YulStmt.expr (YulExpr.call "sstore" [
-                    compatSlotExpr slot,
-                    YulExpr.ident "__compat_value"
-                  ]))
-              )
-            ]
+                    writeSlot,
+                    valueExpr
+                  ])
+                ]
+          | _ =>
+              let compatSlotExpr := fun (slot : Nat) =>
+                let mappingBase := YulExpr.call "mappingSlot" [YulExpr.lit slot, YulExpr.ident "__compat_key"]
+                if wordOffset == 0 then mappingBase else YulExpr.call "add" [mappingBase, YulExpr.lit wordOffset]
+              pure [
+                YulStmt.block (
+                  [YulStmt.let_ "__compat_key" keyExpr, YulStmt.let_ "__compat_value" valueExpr] ++
+                  slots.map (fun slot =>
+                    YulStmt.expr (YulExpr.call storeBuiltin [
+                      compatSlotExpr slot,
+                      YulExpr.ident "__compat_value"
+                    ]))
+                )
+              ]
     | none => throw s!"Compilation error: unknown mapping field '{field}' in {label}"
 
 def compileMappingPackedSlotWrite (fields : List Field) (field : String)
     (keyExpr valueExpr : YulExpr) (wordOffset : Nat) (packed : PackedBits)
-    (label : String) : Except String (List YulStmt) :=
+    (label : String) (allowTransient : Bool := false) : Except String (List YulStmt) :=
   if !isMapping fields field then
     throw s!"Compilation error: field '{field}' is not a mapping"
   else if !packedBitsValid packed then
     throw s!"Compilation error: {label} for field '{field}' has invalid packed range offset={packed.offset} width={packed.width}. Require 0 < width <= 256, offset < 256, and offset + width <= 256."
   else
+    let loadBuiltin :=
+      if allowTransient then
+        match findFieldWithResolvedSlot fields field with
+        | some (f, _) => if f.isTransient then "tload" else "sload"
+        | none => "sload"
+      else
+        "sload"
+    let storeBuiltin :=
+      if allowTransient then
+        match findFieldWithResolvedSlot fields field with
+        | some (f, _) => if f.isTransient then "tstore" else "sstore"
+        | none => "sstore"
+      else
+        "sstore"
     match findFieldWriteSlots fields field with
     | some slots =>
         match slots with
@@ -63,12 +96,12 @@ def compileMappingPackedSlotWrite (fields : List Field) (field : String)
               YulStmt.block [
                 YulStmt.let_ "__compat_value" valueExpr,
                 YulStmt.let_ "__compat_packed" (YulExpr.call "and" [YulExpr.ident "__compat_value", YulExpr.lit maskNat]),
-                YulStmt.let_ "__compat_slot_word" (YulExpr.call "sload" [writeSlot]),
+                YulStmt.let_ "__compat_slot_word" (YulExpr.call loadBuiltin [writeSlot]),
                 YulStmt.let_ "__compat_slot_cleared" (YulExpr.call "and" [
                   YulExpr.ident "__compat_slot_word",
                   YulExpr.call "not" [YulExpr.lit shiftedMaskNat]
                 ]),
-                YulStmt.expr (YulExpr.call "sstore" [
+                YulStmt.expr (YulExpr.call storeBuiltin [
                   writeSlot,
                   YulExpr.call "or" [
                     YulExpr.ident "__compat_slot_cleared",
@@ -90,12 +123,12 @@ def compileMappingPackedSlotWrite (fields : List Field) (field : String)
                  YulStmt.let_ "__compat_packed" (YulExpr.call "and" [YulExpr.ident "__compat_value", YulExpr.lit maskNat])] ++
                 slots.map (fun slot =>
                   YulStmt.block [
-                    YulStmt.let_ "__compat_slot_word" (YulExpr.call "sload" [slotExpr slot]),
+                    YulStmt.let_ "__compat_slot_word" (YulExpr.call loadBuiltin [slotExpr slot]),
                     YulStmt.let_ "__compat_slot_cleared" (YulExpr.call "and" [
                       YulExpr.ident "__compat_slot_word",
                       YulExpr.call "not" [YulExpr.lit shiftedMaskNat]
                     ]),
-                    YulStmt.expr (YulExpr.call "sstore" [
+                    YulStmt.expr (YulExpr.call storeBuiltin [
                       slotExpr slot,
                       YulExpr.call "or" [
                         YulExpr.ident "__compat_slot_cleared",