Make asm constant substitution scope- and flow-aware (bug)

run_contests.sh

@@ -8,6 +8,7 @@ cd "$root_dir"
 bash ./contest.sh test/examples/dispatch/basic.json
 bash ./contest.sh test/examples/dispatch/assembly.json
 bash ./contest.sh test/examples/dispatch/asm_break_continue_leave.json
+bash ./contest.sh test/examples/dispatch/asm_subst.json
 bash ./contest.sh test/examples/dispatch/neg.json
 bash ./contest.sh test/examples/dispatch/miniERC20.json
 bash ./contest.sh test/examples/dispatch/Revert.json

src/Solcore/Backend/MastEval.hs

@@ -16,6 +16,7 @@ module Solcore.Backend.MastEval
     evalYulOp,
     evalYulStmt,
     evalYulBlock,
+    substYulBlock,
     asmIsInterpretable,
     maskWord,
     mstoreBytes,
@@ -707,36 +708,138 @@ venvToSubst env =
     yulLit (StrLit s) = YLit (YulString s)
 
 -- | Substitute known literal values into a Yul block.
--- Only replaces YIdent occurrences in expression positions; does not touch
--- variable names on the left-hand side of let/assign or function parameter lists.
+--
+-- The substitution map is inlined into expression positions only (never into
+-- the names on the left of let/assign or into parameter lists), but it is
+-- threaded through the statement sequence so it stays both scope- and
+-- flow-aware:
+--
+--   * An asm-local @let x@ shadows any outer @x@: @x@ is dropped from the map
+--     for the rest of that scope, so later reads resolve to the local binding
+--     instead of the eliminated comptime value (otherwise a @for@ loop that
+--     rebinds its counter would read the stale outer literal forever).
+--   * An assignment @x := …@ makes the pre-block literal of @x@ stale: @x@ is
+--     dropped from the map so subsequent statements read the new runtime value
+--     (otherwise @x := 5; y := add(x, 1)@ would inline the pre-block @x@).
+--   * @let@/parameter/return names introduced by a nested @if@/@for@/@switch@/
+--     function shadow within their own scope only; assignments to outer
+--     variables inside them escape and invalidate the outer literal afterwards.
 substYulBlock :: Map.Map Name YulExp -> [YulStmt] -> [YulStmt]
-substYulBlock subst = map (substYulStmt subst)
-
-substYulStmt :: Map.Map Name YulExp -> YulStmt -> YulStmt
-substYulStmt subst (YAssign names e) = YAssign names (substYulExp subst e)
-substYulStmt subst (YExp e) = YExp (substYulExp subst e)
-substYulStmt subst (YLet names me) = YLet names (fmap (substYulExp subst) me)
-substYulStmt subst (YIf e block) = YIf (substYulExp subst e) (substYulBlock subst block)
-substYulStmt subst (YBlock stmts) = YBlock (substYulBlock subst stmts)
-substYulStmt subst (YFun n as rs body) = YFun n as rs (substYulBlock subst body)
+substYulBlock _ [] = []
+substYulBlock subst (s : ss) =
+  let (s', subst') = substYulStmt subst s
+   in s' : substYulBlock subst' ss
+
+-- | Substitute into one statement and return the map to use for the statements
+-- that follow it in the same block (with shadowed/reassigned names removed).
+substYulStmt :: Map.Map Name YulExp -> YulStmt -> (YulStmt, Map.Map Name YulExp)
+substYulStmt subst (YAssign names e) =
+  (YAssign names (substYulExp subst e), dropNames names subst)
+substYulStmt subst (YExp e) = (YExp (substYulExp subst e), subst)
+substYulStmt subst (YLet names me) =
+  (YLet names (fmap (substYulExp subst) me), dropNames names subst)
+substYulStmt subst (YIf e block) =
+  ( YIf (substYulExp subst e) (substYulBlock subst block),
+    dropSet (escapingAssignsYulBlock block) subst
+  )
+substYulStmt subst (YBlock stmts) =
+  ( YBlock (substYulBlock subst stmts),
+    dropSet (escapingAssignsYulBlock stmts) subst
+  )
+substYulStmt subst (YFun n as rs body) =
+  -- Parameters and return variables are local to the body; the definition
+  -- itself does not change any binding in the enclosing scope.
+  let inner = dropNames (as ++ concat rs) subst
+   in (YFun n as rs (substYulBlock inner body), subst)
 substYulStmt subst (YFor pre c post b) =
-  YFor
-    (substYulBlock subst pre)
-    (substYulExp subst c)
-    (substYulBlock subst post)
-    (substYulBlock subst b)
+  -- Top-level @pre@ lets scope over the whole loop and shadow outer bindings;
+  -- any variable assigned inside the loop is not loop-invariant, so its
+  -- pre-loop literal must not be inlined into the condition/post/body either.
+  let shadowed = letBoundYulBlock pre
+      escaping =
+        escapingAssignsYulBlock pre
+          `Set.union` escapingAssignsYulBlock post
+          `Set.union` escapingAssignsYulBlock b
+      inner = dropSet (shadowed `Set.union` escaping) subst
+   in ( YFor
+          (substYulBlock inner pre)
+          (substYulExp inner c)
+          (substYulBlock inner post)
+          (substYulBlock inner b),
+        dropSet escaping subst
+      )
 substYulStmt subst (YSwitch e cases def) =
-  YSwitch
-    (substYulExp subst e)
-    (map (\(lit, block) -> (lit, substYulBlock subst block)) cases)
-    (fmap (substYulBlock subst) def)
-substYulStmt _ stmt = stmt -- YBreak, YContinue, YLeave, YComment unchanged
+  let cases' = map (\(lit, block) -> (lit, substYulBlock subst block)) cases
+      def' = fmap (substYulBlock subst) def
+      escaping =
+        foldMap (escapingAssignsYulBlock . snd) cases
+          `Set.union` maybe Set.empty escapingAssignsYulBlock def
+   in ( YSwitch (substYulExp subst e) cases' def',
+        dropSet escaping subst
+      )
+substYulStmt subst stmt = (stmt, subst) -- YBreak, YContinue, YLeave, YComment
+
+dropNames :: [Name] -> Map.Map Name YulExp -> Map.Map Name YulExp
+dropNames names m = foldr Map.delete m names
+
+dropSet :: Set.Set Name -> Map.Map Name YulExp -> Map.Map Name YulExp
+dropSet names m = Set.foldr Map.delete m names
 
 substYulExp :: Map.Map Name YulExp -> YulExp -> YulExp
 substYulExp subst (YIdent n) = Map.findWithDefault (YIdent n) n subst
 substYulExp subst (YCall op args) = YCall op (map (substYulExp subst) args)
 substYulExp _ e = e -- YLit, YMeta unchanged
 
+-- | Names @let@-bound at the top level of a block; these are the only ones a
+-- @for@ pre-block puts in scope for the rest of the loop.
+letBoundYulBlock :: [YulStmt] -> Set.Set Name
+letBoundYulBlock = foldMap letBound
+  where
+    letBound (YLet names _) = Set.fromList names
+    letBound _ = Set.empty
+
+-- | Names assigned inside a block that refer to a variable declared outside it,
+-- i.e. assignments whose effect escapes the block's own scope. A name that is
+-- @let@-bound anywhere within the block is local, so it is excluded — dropping
+-- an outer literal for a purely local reassignment would be both unnecessary
+-- and unsafe (the outer name may be an eliminated comptime value still read
+-- after the block). Yul functions have an isolated scope, so their bodies are
+-- not traversed.
+escapingAssignsYulBlock :: [YulStmt] -> Set.Set Name
+escapingAssignsYulBlock block =
+  assignedYulBlock block `Set.difference` letBoundDeepYulBlock block
+
+assignedYulBlock :: [YulStmt] -> Set.Set Name
+assignedYulBlock = foldMap assignedYulStmt
+
+assignedYulStmt :: YulStmt -> Set.Set Name
+assignedYulStmt (YAssign names _) = Set.fromList names
+assignedYulStmt (YIf _ block) = assignedYulBlock block
+assignedYulStmt (YBlock stmts) = assignedYulBlock stmts
+assignedYulStmt (YFor pre _ post b) =
+  assignedYulBlock pre `Set.union` assignedYulBlock post `Set.union` assignedYulBlock b
+assignedYulStmt (YSwitch _ cases def) =
+  foldMap (assignedYulBlock . snd) cases
+    `Set.union` maybe Set.empty assignedYulBlock def
+assignedYulStmt (YFun {}) = Set.empty -- isolated scope
+assignedYulStmt _ = Set.empty
+
+letBoundDeepYulBlock :: [YulStmt] -> Set.Set Name
+letBoundDeepYulBlock = foldMap go
+  where
+    go (YLet names _) = Set.fromList names
+    go (YIf _ block) = letBoundDeepYulBlock block
+    go (YBlock stmts) = letBoundDeepYulBlock stmts
+    go (YFor pre _ post b) =
+      letBoundDeepYulBlock pre
+        `Set.union` letBoundDeepYulBlock post
+        `Set.union` letBoundDeepYulBlock b
+    go (YSwitch _ cases def) =
+      foldMap (letBoundDeepYulBlock . snd) cases
+        `Set.union` maybe Set.empty letBoundDeepYulBlock def
+    go (YFun {}) = Set.empty -- isolated scope
+    go _ = Set.empty
+
 -- | Merge a YulState back into VEnv, using TypeReg to find the right MastIds.
 -- Only names present in TypeReg are merged; others are silently ignored.
 mergeYulStateToVEnv :: TypeReg -> YulState -> VEnv -> VEnv

test/Cases.hs

@@ -128,6 +128,7 @@ dispatches =
     "Files for dispatch cases"
     [ runDispatchTest "basic.solc",
       runDispatchTest "assembly.solc",
+      runDispatchTest "asm_subst.solc",
       runDispatchTest "stringid.solc",
       runDispatchTest "storage.solc",
       runDispatchTest "miniERC20.solc",

test/YulEvalTests.hs

@@ -30,6 +30,13 @@ yExp op args = YExp (YCall (Name op) args)
 st :: [(String, Integer)] -> YulState
 st = Map.fromList . map (\(n, v) -> (Name n, v))
 
+-- A name→literal substitution map, as built from a VEnv before an asm block.
+mkSubst :: [(String, Integer)] -> Map.Map Name YulExp
+mkSubst = Map.fromList . map (\(n, v) -> (Name n, YLit (YulNumber v)))
+
+yLet :: String -> YulExp -> YulStmt
+yLet n e = YLet [Name n] (Just e)
+
 -- Run an EvalM action in non-comptime mode (memory ops inactive).
 runPure :: EvalM a -> a
 runPure m = fst $ runEvalM (EvalEnv Map.empty Set.empty False) defaultFuel m
@@ -56,6 +63,7 @@ yulEvalTests =
       evalYulOpTests,
       evalYulExpTests,
       evalYulBlockTests,
+      substYulBlockTests,
       memoryHelperTests,
       memoryEvalTests,
       asmIsInterpretableTests
@@ -270,6 +278,93 @@ evalYulBlockTests =
           @?= Nothing
     ]
 
+-----------------------------------------------------------------------
+-- substYulBlock: scope- and flow-aware inlining of comptime literals
+-----------------------------------------------------------------------
+
+substYulBlockTests :: TestTree
+substYulBlockTests =
+  testGroup
+    "substYulBlock (scope/flow-aware substitution)"
+    [ testCase "an unassigned variable is still inlined" $
+        substYulBlock
+          (mkSubst [("x", 0)])
+          [yAssign "z" (yCall "add" [yIdent "x", yNum 1])]
+          @?= [yAssign "z" (yCall "add" [yNum 0, yNum 1])],
+      -- Flow: after `x := 5`, `x` holds a new runtime value, so the pre-block
+      -- literal must not be inlined into the following `add(x, 1)`.
+      testCase "a reassigned variable is not inlined downstream" $
+        let block =
+              [ yAssign "x" (yNum 5),
+                yAssign "y" (yCall "add" [yIdent "x", yNum 1])
+              ]
+         in substYulBlock (mkSubst [("x", 0), ("y", 0)]) block @?= block,
+      -- Shadowing: the asm-local `let i` owns the name for the rest of the
+      -- block, so `add(i, 1)` must read the local, not the outer literal.
+      testCase "a plain asm-local let shadows the outer literal" $
+        let block =
+              [ yLet "i" (yNum 7),
+                yAssign "r" (yCall "add" [yIdent "i", yNum 1])
+              ]
+         in substYulBlock (mkSubst [("i", 0)]) block @?= block,
+      -- ...but the let's initialiser is still evaluated in the outer scope.
+      testCase "a let initialiser is inlined from the outer scope" $
+        substYulBlock
+          (mkSubst [("x", 9)])
+          [yLet "i" (yIdent "x"), yAssign "r" (yIdent "i")]
+          @?= [yLet "i" (yNum 9), yAssign "r" (yIdent "i")],
+      -- The reported infinite-loop case: the for-counter is `let`-bound in the
+      -- pre-block, so neither the condition nor the body may inline the outer 0.
+      testCase "a for-loop counter shadowed by its pre-let is not inlined" $
+        let block =
+              [ YFor
+                  [yLet "i" (yNum 0)]
+                  (yCall "lt" [yIdent "i", yNum 3])
+                  [yAssign "i" (yCall "add" [yIdent "i", yNum 1])]
+                  [yAssign "r" (yIdent "i")]
+              ]
+         in substYulBlock (mkSubst [("i", 0)]) block @?= block,
+      -- A variable assigned inside the loop is not loop-invariant either.
+      testCase "a variable assigned in a for-body is not inlined into it" $
+        let block =
+              [ YFor
+                  []
+                  (yCall "lt" [yIdent "n", yNum 3])
+                  []
+                  [yAssign "s" (yCall "add" [yIdent "s", yNum 1])]
+              ]
+         in substYulBlock (mkSubst [("s", 0)]) block @?= block,
+      -- Function parameters and returns shadow within the body.
+      testCase "function parameters and returns shadow within the body" $
+        let block =
+              [ YFun
+                  (Name "f")
+                  [Name "a"]
+                  (Just [Name "r"])
+                  [yAssign "r" (yCall "add" [yIdent "a", yNum 1])]
+              ]
+         in substYulBlock (mkSubst [("a", 0), ("r", 0)]) block @?= block,
+      -- ...but a function definition does not disturb the enclosing scope: a
+      -- statement after it still sees the outer binding.
+      testCase "a function definition does not leak its shadowing outward" $
+        substYulBlock
+          (mkSubst [("a", 0)])
+          [ YFun (Name "f") [Name "a"] (Just [Name "r"]) [yAssign "r" (yIdent "a")],
+            yAssign "z" (yIdent "a")
+          ]
+          @?= [ YFun (Name "f") [Name "a"] (Just [Name "r"]) [yAssign "r" (yIdent "a")],
+                yAssign "z" (yNum 0)
+              ],
+      -- An assignment to an outer variable inside a nested block escapes and
+      -- invalidates the literal for statements after the block.
+      testCase "an assignment inside an if escapes the block" $
+        let block =
+              [ YIf (yIdent "flag") [yAssign "x" (yNum 5)],
+                yAssign "y" (yIdent "x")
+              ]
+         in substYulBlock (mkSubst [("x", 0)]) block @?= block
+    ]
+
 -----------------------------------------------------------------------
 -- Memory helpers (pure functions: mstoreBytes, mloadWord)
 -----------------------------------------------------------------------

test/examples/dispatch/asm_subst.json

@@ -0,0 +1,40 @@
+{
+  "asm_subst": {
+    "bytecode": "_CODE",
+    "contract": "C",
+    "tests": [
+      {
+        "input": {
+          "comment": "constructor()",
+          "calldata": "",
+          "value": "0"
+        },
+        "kind": "constructor"
+      },
+      {
+        "input": {
+          "comment": "asmFlow()(uint256) -> 6 (x := 5; y := add(x, 1))",
+          "calldata": "8decb8c5",
+          "value": "0"
+        },
+        "kind": "call",
+        "output": {
+          "returndata": "0000000000000000000000000000000000000000000000000000000000000006",
+          "status": "success"
+        }
+      },
+      {
+        "input": {
+          "comment": "asmLoopFlow()(uint256) -> 3 (accumulate s across 3 iterations)",
+          "calldata": "2b591092",
+          "value": "0"
+        },
+        "kind": "call",
+        "output": {
+          "returndata": "0000000000000000000000000000000000000000000000000000000000000003",
+          "status": "success"
+        }
+      }
+    ]
+  }
+}

test/examples/dispatch/asm_subst.solc

@@ -0,0 +1,40 @@
+import std.{*};
+import std.dispatch.{*};
+
+// Regression tests for constant substitution into `assembly` blocks, which the
+// partial evaluator must apply in a scope- and flow-aware way.  A single,
+// uniform substitution map (built from every literal-valued local in scope)
+// used to be inlined into the whole block at once, producing wrong code.
+contract C {
+    constructor() {}
+
+    // Flow-sensitivity: after `x := 5` the read of `x` in `add(x, 1)` must see
+    // the new value 5, not the pre-block literal 0.  The buggy substitution
+    // folded the second statement to `add(0, 1)` and returned 1; the correct
+    // result is 6.
+    public function asmFlow() -> uint256 {
+        let x : word = 0;
+        let y : word = 0;
+        assembly {
+          x := 5
+          y := add(x, 1)
+        }
+        return uint256(y);
+    }
+
+    // A variable assigned inside a loop is not loop-invariant: the pre-loop
+    // literal `s == 0` must not be inlined into `add(s, 1)` in the body.  The
+    // buggy substitution rewrote the body to `s := add(0, 1)`, so `s` was reset
+    // to 1 on every iteration and the function returned 1; the correct sum over
+    // three iterations is 3.  The asm-local counter `j` (no outer counterpart)
+    // exercises the loop without any name collision.
+    public function asmLoopFlow() -> uint256 {
+        let s : word = 0;
+        assembly {
+          for { let j := 0 } lt(j, 3) { j := add(j, 1) } {
+            s := add(s, 1)
+          }
+        }
+        return uint256(s);
+    }
+}