Support dynamic/composite internal helper args (#1889)

Compiler/CompilationModel.lean

@@ -13,6 +13,7 @@ import Compiler.CompilationModel.DynamicData
 import Compiler.CompilationModel.EcmAxiomCollection
 import Compiler.CompilationModel.EventEmission
 import Compiler.CompilationModel.EventAbiHelpers
+import Compiler.CompilationModel.InternalArgs
 import Compiler.CompilationModel.InternalNaming
 import Compiler.CompilationModel.IssueRefs
 import Compiler.CompilationModel.LayoutReport

Compiler/CompilationModel/Dispatch.lean

@@ -5,6 +5,7 @@
   the lower-level statement/expression compilation helpers.
 -/
 import Compiler.CompilationModel.Compile
+import Compiler.CompilationModel.InternalArgs
 import Compiler.CompilationModel.ParamLoading
 import Compiler.CompilationModel.ScopeValidation
 import Compiler.CompilationModel.TrustSurface
@@ -35,41 +36,18 @@ def freshInternalRetNames (returns : List ParamType) (usedNames : List String) :
     (usedNames, [])
   namesRev.reverse
 
-def internalFunctionYulParamNames (params : List Param) : List String :=
-  params.flatMap fun param =>
-    match param.ty with
-    | ParamType.array _ =>
-        [s!"{param.name}_data_offset", s!"{param.name}_length"]
-    | ParamType.bytes | ParamType.string =>
-        [s!"{param.name}_data_offset", s!"{param.name}_length"]
-    | ParamType.fixedArray _ _ =>
-        if isDynamicParamType param.ty then
-          [s!"{param.name}_data_offset"]
-        else
-          staticParamBindingNames param.name param.ty
-    | ParamType.tuple _ =>
-        if isDynamicParamType param.ty then
-          [s!"{param.name}_data_offset"]
-        else
-          staticParamBindingNames param.name param.ty
-    | ParamType.newtypeOf _ baseTy =>
-        if isDynamicParamType param.ty then
-          [s!"{param.name}_data_offset"]
-        else
-          staticParamBindingNames param.name baseTy
-    | _ => [param.name]
-
 -- Compile internal function to a Yul function definition (#181)
 def compileInternalFunction (fields : List Field) (events : List EventDef) (errors : List ErrorDef)
-    (adtTypes : List AdtTypeDef := []) (spec : FunctionSpec) :
+    (adtTypes : List AdtTypeDef := []) (spec : FunctionSpec)
+    (internalFunctions : List FunctionSpec := []) :
     Except String YulStmt := do
   validateFunctionSpec spec
   let returns ← functionReturns spec
   let paramNames := internalFunctionYulParamNames spec.params
   let usedNames := paramNames ++ collectStmtListBindNames spec.body
   let retNames := freshInternalRetNames returns usedNames
   let bodyStmts ← compileStmtList fields events errors .calldata retNames true
-    (paramNames ++ retNames) adtTypes spec.body
+    (paramNames ++ retNames) adtTypes spec.body internalFunctions
   pure (YulStmt.funcDef (internalFunctionYulName spec.name) paramNames retNames bodyStmts)
 
 theorem compileInternalFunction_ok_components
@@ -189,13 +167,14 @@ theorem compileInternalFunction_some_ok_of_components
 
 -- Compile function spec to IR function
 def compileFunctionSpec (fields : List Field) (events : List EventDef) (errors : List ErrorDef)
-    (adtTypes : List AdtTypeDef := []) (selector : Nat) (spec : FunctionSpec) :
+    (adtTypes : List AdtTypeDef := []) (selector : Nat) (spec : FunctionSpec)
+    (internalFunctions : List FunctionSpec := []) :
     Except String IRFunction := do
   validateFunctionSpec spec
   let returns ← functionReturns spec
   let paramLoads := genParamLoads spec.params
   let bodyStmts ← compileStmtList fields events errors .calldata [] false
-    (spec.params.map (·.name)) adtTypes spec.body
+    (spec.params.map (·.name)) adtTypes spec.body internalFunctions
   let allStmts := paramLoads ++ bodyStmts
   let retType := match returns with
     | [single] => single.toIRType
@@ -276,9 +255,10 @@ def attachNonReentrantGuard (fields : List Field) (spec : FunctionSpec)
       pure { irFn with body := prefixLoads ++ guardStmts ++ suffix ++ [release] }
 
 private def compileSpecialEntrypoint (fields : List Field) (events : List EventDef)
-    (errors : List ErrorDef) (adtTypes : List AdtTypeDef := []) (spec : FunctionSpec) :
+    (errors : List ErrorDef) (adtTypes : List AdtTypeDef := [])
+    (internalFunctions : List FunctionSpec := []) (spec : FunctionSpec) :
     Except String IREntrypoint := do
-  let bodyChunks ← compileStmtList fields events errors .calldata [] false [] adtTypes spec.body
+  let bodyChunks ← compileStmtList fields events errors .calldata [] false [] adtTypes spec.body internalFunctions
   -- Apply nonreentrant guard for fallback/receive if annotated (high-severity
   -- Bugbot: previously these special entrypoints were compiled without the
   -- transient lock even when `nonreentrant(lock)` was declared).
@@ -310,14 +290,15 @@ def usesMapping (fields : List Field) : Bool :=
 -- Compile deploy code (constructor)
 -- Note: Don't append datacopy/return here - Codegen.deployCode does that
 def compileConstructor (fields : List Field) (events : List EventDef) (errors : List ErrorDef)
-    (adtTypes : List AdtTypeDef := []) (ctor : Option ConstructorSpec) :
+    (adtTypes : List AdtTypeDef := []) (ctor : Option ConstructorSpec)
+    (internalFunctions : List FunctionSpec := []) :
     Except String (List YulStmt) := do
   match ctor with
   | none => return []
   | some spec =>
     let argLoads := genConstructorArgLoads spec.params
     let bodyChunks ← compileStmtList fields events errors .memory [] false
-      (spec.params.map (·.name)) adtTypes spec.body
+      (spec.params.map (·.name)) adtTypes spec.body internalFunctions
     return argLoads ++ bodyChunks
 
 -- Main compilation function
@@ -412,7 +393,7 @@ private def validateCompileInputsBeforeFieldWriteConflict
   | some ctor => do
       ctor.body.forM (validateEventArgShapesInStmt "constructor" ctor.params spec.events)
       ctor.body.forM (validateCustomErrorArgShapesInStmt "constructor" ctor.params spec.errors)
-      ctor.body.forM (validateInternalCallShapesInStmt spec.functions "constructor")
+      ctor.body.forM (validateInternalCallShapesInStmt spec.functions "constructor" ctor.params)
   for ext in spec.externals do
     let _ ← externalFunctionReturns ext
     validateInteropExternalSpec ext
@@ -552,8 +533,9 @@ def validateCompileInputs (spec : CompilationModel) (selectors : List Nat)
     `compileFunctionSpec` (see `attachNonReentrantGuard`). -/
 def compileGuardedFunctionSpec (fields : List Field) (events : List EventDef)
     (errors : List ErrorDef) (adtTypes : List AdtTypeDef)
+    (internalFunctions : List FunctionSpec)
     (sel : Nat) (fnSpec : FunctionSpec) : Except String IRFunction := do
-  let irFn ← compileFunctionSpec fields events errors adtTypes sel fnSpec
+  let irFn ← compileFunctionSpec fields events errors adtTypes sel fnSpec internalFunctions
   attachNonReentrantGuard fields fnSpec irFn
 
 def compileValidatedCore (spec : CompilationModel) (selectors : List Nat) : Except String IRContract := do
@@ -570,8 +552,9 @@ def compileValidatedCore (spec : CompilationModel) (selectors : List Nat) : Exce
   let fallbackSpec ← pickUniqueFunctionByName "fallback" spec.functions
   let receiveSpec ← pickUniqueFunctionByName "receive" spec.functions
   let functions ← (externalFns.zip selectors).mapM fun entry =>
-    compileGuardedFunctionSpec fields spec.events spec.errors spec.adtTypes entry.2 entry.1
-  let internalFuncDefs ← internalFns.mapM (compileInternalFunction fields spec.events spec.errors spec.adtTypes)
+    compileGuardedFunctionSpec fields spec.events spec.errors spec.adtTypes internalFns entry.2 entry.1
+  let internalFuncDefs ← internalFns.mapM fun fn =>
+    compileInternalFunction fields spec.events spec.errors spec.adtTypes fn internalFns
   let arrayElementHelpers :=
     (if arrayHelpersRequired then
       [ checkedArrayElementCalldataHelper
@@ -631,11 +614,11 @@ def compileValidatedCore (spec : CompilationModel) (selectors : List Nat) : Exce
       [dynamicBytesEqCalldataHelper, dynamicBytesEqMemoryHelper]
     else
       []
-  let fallbackEntrypoint ← fallbackSpec.mapM (compileSpecialEntrypoint fields spec.events spec.errors spec.adtTypes)
-  let receiveEntrypoint ← receiveSpec.mapM (compileSpecialEntrypoint fields spec.events spec.errors spec.adtTypes)
+  let fallbackEntrypoint ← fallbackSpec.mapM (compileSpecialEntrypoint fields spec.events spec.errors spec.adtTypes internalFns)
+  let receiveEntrypoint ← receiveSpec.mapM (compileSpecialEntrypoint fields spec.events spec.errors spec.adtTypes internalFns)
   return {
     name := spec.name
-    deploy := (← compileConstructor fields spec.events spec.errors spec.adtTypes spec.constructor)
+    deploy := (← compileConstructor fields spec.events spec.errors spec.adtTypes spec.constructor internalFns)
     constructorPayable := spec.constructor.map (·.isPayable) |>.getD false
     functions := functions
     fallbackEntrypoint := fallbackEntrypoint

Compiler/CompilationModel/EventEmission.lean

@@ -55,7 +55,8 @@ structure EventDynamicArraySource where
   source : DynamicDataSource
 
 def eventDynamicArraySource?
-    (fields : List Field) (dynamicSource : DynamicDataSource) :
+    (fields : List Field) (dynamicSource : DynamicDataSource)
+    (internalFunctions : List FunctionSpec := []) :
     Expr → Except String (Option EventDynamicArraySource)
   | Expr.param name =>
       pure (some
@@ -68,14 +69,14 @@ def eventDynamicArraySource?
           dataOffsetExpr := YulExpr.ident s!"{name}_data_offset"
           source := .memory })
   | e@(Expr.paramDynamicMemberLength name wordOffset) => do
-      let dataOffsetExpr ← compileExpr fields dynamicSource
+      let dataOffsetExpr ← compileExprWithInternals fields dynamicSource internalFunctions
         (Expr.paramDynamicMemberDataOffset name wordOffset)
-      let lengthExpr ← compileExpr fields dynamicSource e
+      let lengthExpr ← compileExprWithInternals fields dynamicSource internalFunctions e
       pure (some { lengthExpr, dataOffsetExpr, source := dynamicSource })
   | e@(Expr.arrayElementDynamicMemberLength name index wordOffset) => do
-      let dataOffsetExpr ← compileExpr fields dynamicSource
+      let dataOffsetExpr ← compileExprWithInternals fields dynamicSource internalFunctions
         (Expr.arrayElementDynamicMemberDataOffset name index wordOffset)
-      let lengthExpr ← compileExpr fields dynamicSource e
+      let lengthExpr ← compileExprWithInternals fields dynamicSource internalFunctions e
       pure (some { lengthExpr, dataOffsetExpr, source := dynamicSource })
   | _ => pure none
 
@@ -187,15 +188,16 @@ def compileScalarEmitFromCompiledArgs
 
 def compileEmit (fields : List Field) (events : List EventDef)
     (dynamicSource : DynamicDataSource := .calldata)
-    (eventName : String) (args : List Expr) : Except String (List YulStmt) := do
+    (eventName : String) (args : List Expr) (internalFunctions : List FunctionSpec := []) :
+    Except String (List YulStmt) := do
   let eventDef? := events.find? (·.name == eventName)
   let eventDef ←
     match eventDef? with
     | some e => pure e
     | none => throw s!"Compilation error: unknown event '{eventName}'"
   if args.length != eventDef.params.length then
     throw s!"Compilation error: event '{eventName}' expects {eventDef.params.length} args, got {args.length}"
-  let compiledArgs ← compileExprList fields dynamicSource args
+  let compiledArgs ← compileExprListWithInternals fields dynamicSource internalFunctions args
   let zippedWithSource := eventZippedWithSource eventDef args compiledArgs
   let indexed := eventIndexedArgs zippedWithSource
   let unindexed := eventUnindexedArgs zippedWithSource
@@ -344,7 +346,7 @@ def compileEmit (fields : List Field) (events : List EventDef)
                   | _ =>
                       throw s!"Compilation error: unindexed dynamic array event param '{p.name}' in event '{eventName}' currently requires direct parameter reference ({issue586Ref})."
               else if indexedDynamicArrayElemSupported elemTy then
-                match ← eventDynamicArraySource? fields dynamicSource srcExpr with
+                match ← eventDynamicArraySource? fields dynamicSource internalFunctions srcExpr with
                 | some source =>
                     let lenName := s!"__evt_arg{argIdx}_len"
                     let dstName := s!"__evt_arg{argIdx}_dst"
@@ -527,7 +529,7 @@ def compileEmit (fields : List Field) (events : List EventDef)
                 throw s!"Compilation error: indexed dynamic array event param '{p.name}' in event '{eventName}' currently requires direct parameter reference ({issue586Ref})."
         | _ =>
             if indexedDynamicArrayElemSupported elemTy then
-              match ← eventDynamicArraySource? fields dynamicSource srcExpr with
+              match ← eventDynamicArraySource? fields dynamicSource internalFunctions srcExpr with
               | some source =>
                   let topicName := s!"__evt_topic{idx + 1}"
                   let byteLenName := s!"__evt_arg{idx}_byte_len"