Argument inlining

src/analyzer.fs

@@ -218,9 +218,9 @@ let resolve topLevel =
 
     let resolveGlobalsAndParameters = function
         | TLDecl decl -> resolveDecl VarScope.Global decl
-        | Function (funcType, _) ->
+        | Function (funcType, _) as tl ->
             for decl in funcType.args do resolveDecl VarScope.Parameter decl
-            funcType.fName.Declaration <- Declaration.Func (new FunDecl(funcType))
+            funcType.fName.Declaration <- Declaration.Func (new FunDecl(tl, funcType))
         | _ -> ()
 
     // First visit all declarations, creating them.
@@ -236,39 +236,44 @@ type CallSite = {
     ident: Ident
     varsInScope: string list
     prototype: string * int
+    argExprs: Expr list
 }
 type FuncInfo = {
     func: TopLevel
     funcType: FunctionType
     body: Stmt
     name: string
     callSites: CallSite list // calls to other user-defined functions, from inside this function.
+    isResolvable: bool // Currently we cannot resolve overloaded functions based on argument types.
+    isOverloaded: bool
 }
 let findFuncInfos code =
     let findCallSites block = // Gets the list of call sites in this function
         let callSites = List()
         let collect (mEnv : MapEnv) = function
             | FunCall (Var id, argExprs) as e ->
-                callSites.Add { ident = id; varsInScope = mEnv.vars.Keys |> Seq.toList; prototype = (id.Name, argExprs.Length) }
+                callSites.Add { ident = id; varsInScope = mEnv.vars.Keys |> Seq.toList; prototype = (id.Name, argExprs.Length); argExprs = argExprs }
                 e
             | e -> e
         mapStmt (mapEnv collect id) block |> ignore<MapEnv * Stmt>
         callSites |> Seq.toList
     let functions = code |> List.choose (function
         | Function(funcType, block) as f -> Some (funcType, funcType.fName.Name, block, f)
         | _ -> None)
-    let funcInfos = functions |> List.map (fun (funcType, name, block, f) ->
+    let funcInfos = functions |> List.map (fun ((funcType, name, block, func) as f) ->
         let callSites = findCallSites block
                         // only return calls to user-defined functions
-                        |> List.filter (fun callSite -> functions |> List.exists (fun (ft,_,_,_) -> callSite.prototype = ft.prototype))
-        { FuncInfo.func = f; funcType = funcType; name = name; callSites = callSites; body = block })
+                        |> List.filter (fun callSite -> functions |> List.exists (fun (ft, _, _, _) -> callSite.prototype = ft.prototype))
+        let isResolvable = not (functions |> List.except [f] |> List.exists (fun (ft, _, _, _) -> ft.prototype = funcType.prototype))
+        let isOverloaded = not (functions |> List.except [f] |> List.exists (fun (ft, _, _, _) -> ft.fName = funcType.fName))
+        { FuncInfo.func = func; funcType = funcType; name = name; callSites = callSites; body = block; isResolvable = isResolvable; isOverloaded = isOverloaded })
     funcInfos
 
 module private FunctionInlining =
 
     // To ensure correctness, we verify if it's safe to inline.
     //
-    // [A] Only inline a function if it never refers to a global by a name function or variable that is shadowed by a local variable in scope at the call site.
+    // [A] Only inline a function if it never refers to a global function or variable by a name that is shadowed by a local variable in scope at the call site.
     // [B] Only inline a function if it has only one call site.
     //     Exception: if the body is "trivial" it will be inlined at all call sites.
     // [C] Only inline a function if it is a single expression return.
@@ -325,17 +330,14 @@ module private FunctionInlining =
         let funcInfos = findFuncInfos code
         for funcInfo in funcInfos do
             let canBeRenamed = not (options.noRenamingList |> List.contains funcInfo.name) // noRenamingList includes "main"
-            let isExternal = options.hlsl && funcInfo.funcType.semantics <> []
-            let isOverloadedAmbiguously = funcInfos |> List.except [funcInfo] |> List.exists (fun f -> f.funcType.prototype = funcInfo.funcType.prototype)
-            if canBeRenamed && not isExternal && not isOverloadedAmbiguously then
+            if canBeRenamed && not funcInfo.funcType.isExternal && funcInfo.isResolvable then
                 if not funcInfo.funcType.hasOutOrInoutParams then // [F]
                     // Find calls to this function. This works because we checked that the function is not overloaded ambiguously.
                     let callSites = funcInfos |> List.collect (fun n -> n.callSites)
                                               |> List.filter (fun callSite -> callSite.prototype = funcInfo.funcType.prototype)
                     if callSites.Length > 0 then // Unused function elimination is not handled here
-                        match funcInfo.body with
-                        | Jump (JumpKeyword.Return, Some body)
-                        | Block [Jump (JumpKeyword.Return, Some body)] -> // [C]
+                        match funcInfo.body.asStmtList with
+                        | [Jump (JumpKeyword.Return, Some body)] -> // [C]
                             if callSites.Length = 1 || VariableInlining.isTrivialExpr body then // [B]
                                 tryMarkFunctionToInline funcInfo callSites
                         | _ -> ()

src/ast.fs

@@ -12,6 +12,7 @@ type Ident(name: string) =
     member this.OldName = name
     member this.Rename(n) = newName <- n
     member val ToBeInlined = newName.StartsWith("i_") with get, set
+    // This prefix disables function inlining and variable inlining.
     member this.DoNotInline = this.OldName.StartsWith("noinline_")
 
     //member val isVarRead: bool = false with get, set
@@ -21,7 +22,7 @@ type Ident(name: string) =
                           | Declaration.Variable rv -> Some rv
                           | _ -> None
 
-     // Real identifiers cannot start with a digit, but the temporary ids of the rename pass are numbers.
+    // Real identifiers cannot start with a digit, but the temporary ids of the rename pass are numbers.
     member this.IsUniqueId = System.Char.IsDigit this.Name.[0]
 
     interface System.IComparable with
@@ -41,12 +42,13 @@ and [<NoComparison>] [<RequireQualifiedAccess>] Declaration =
     | Variable of VarDecl
     | Func of FunDecl
 and VarDecl(ty, decl, scope) =
-    member val ty: Type = ty with get, set
-    member val decl = decl: DeclElt with get, set
-    member val scope = scope: VarScope with get, set
+    member val ty: Type = ty with get
+    member val decl = decl: DeclElt with get
+    member val scope = scope: VarScope with get
     member val isEverWrittenAfterDecl = false with get, set
-and FunDecl(funcType) =
-    member val funcType: FunctionType = funcType with get, set
+and FunDecl(func, funcType) =
+    member val func: TopLevel = func with get
+    member val funcType: FunctionType = funcType with get
     member val hasExternallyVisibleSideEffects = false with get, set
 and [<RequireQualifiedAccess>] JumpKeyword = Break | Continue | Discard | Return
 
@@ -113,6 +115,9 @@ and Stmt =
     | Jump of JumpKeyword * Expr option (*break, continue, return (expr)?, discard*)
     | Verbatim of string
     | Switch of Expr * (CaseLabel * Stmt list) list
+with member this.asStmtList = match this with
+                              | Block stmts -> stmts
+                              | stmt -> [stmt]
 
 and CaseLabel =
     | Case of Expr
@@ -124,10 +129,17 @@ and FunctionType = {
     args: Decl list (*args*)
     semantics: Expr list (*semantics*)
 } with
+    member this.isExternal = options.hlsl && this.semantics <> []
     member this.hasOutOrInoutParams =
         let typeQualifiers = set [for (ty, _) in this.args do yield! ty.typeQ]
         not (Set.intersect typeQualifiers (set ["out"; "inout"])).IsEmpty
     member this.prototype = (this.fName.Name, this.args.Length)
+    member this.parameters: (Type * DeclElt) list =
+        // Helper to extract the single DeclElt for each arg.
+        [ for (ty, declElts) in this.args do
+          yield match declElts with
+                | [declElt] -> ty, declElt
+                | _ -> failwith "invalid declElt for function argument" ]
     override t.ToString() =
         let sem = if t.semantics.IsEmpty then "" else $": {t.semantics}" in
         let args = System.String.Join(", ", t.args |> List.map (function

src/builtin.fs

@@ -45,6 +45,15 @@ let assignOps = set [
     "<<="; ">>="; "&="; "^="; "|="
     "++"; "--"; "$++"; "$--"
 ]
+let nonAssignOps = set [
+    "*"; "/"; "%"
+    "+"; "-"
+    "<<"; ">>"
+    "<"; ">"; "<="; ">="
+    "=="; "!="
+    "&"; "^"; "|"
+    "&&"; "^^"; "||"
+]
 
 let castFunctions = builtinTypes - set ["void"]
 let trigonometryFunctions = set([

src/renamer.fs

@@ -206,8 +206,7 @@ module private RenamerImpl =
             env
 
     let renFctName env (f: FunctionType) =
-        let isExternal = options.hlsl && f.semantics <> []
-        if (isExternal && options.preserveExternals) || options.preserveAllGlobals then
+        if (f.isExternal && options.preserveExternals) || options.preserveAllGlobals then
             env
         elif List.contains f.fName.Name options.noRenamingList then
             env
@@ -218,7 +217,7 @@ module private RenamerImpl =
                 env
             | None ->
                 let newEnv = renFunction env (List.length f.args) f.fName
-                if isExternal then export env ExportPrefix.HlslFunction f.fName
+                if f.isExternal then export env ExportPrefix.HlslFunction f.fName
                 newEnv
 
     let renList env fct li =

src/rewriter.fs

@@ -272,9 +272,8 @@ let private simplifyExpr (didInline: bool ref) env = function
         | Some ([{args = declArgs}, body]) ->
             if List.length declArgs <> List.length passedArgs then
                 failwithf "Cannot inline function %s since it doesn't have the right number of arguments" v.Name
-            match body with
-            | Jump (JumpKeyword.Return, Some bodyExpr)
-            | Block [Jump (JumpKeyword.Return, Some bodyExpr)] ->
+            match body.asStmtList with
+            | [Jump (JumpKeyword.Return, Some bodyExpr)] ->
                 didInline.Value <- true
                 inlineFn declArgs passedArgs bodyExpr
             // Don't yell if we've done some inlining this pass -- maybe it
@@ -556,12 +555,11 @@ let rec private removeUnusedFunctions code =
     let funcInfos = Analyzer.findFuncInfos code
     let isUnused (funcInfo : Analyzer.FuncInfo) =
         let canBeRenamed = not (options.noRenamingList |> List.contains funcInfo.name) // noRenamingList includes "main"
-        let isExternal = options.hlsl && funcInfo.funcType.semantics <> []
         let isCalled = funcInfos |> List.exists (fun n ->
             n.callSites
             |> List.map (fun c -> c.prototype)
             |> List.contains funcInfo.funcType.prototype) // when in doubt wrt overload resolution, keep the function.
-        canBeRenamed && not isCalled && not isExternal
+        canBeRenamed && not isCalled && not funcInfo.funcType.isExternal
     let unused = set [for funcInfo in funcInfos do if isUnused funcInfo then yield funcInfo.func]
     let mutable edited = false
     let code = code |> List.filter (function
@@ -592,6 +590,87 @@ let reorderFunctions code =
     rest @ order
 
 
+// Inline the argument of a function call into the function body.
+module private ArgumentInlining =
+
+    let rec isInlinableExpr = function
+        | Var v when v.Name = "true" || v.Name = "false" -> true
+        | Int _
+        | Float _ -> true
+        | FunCall(Var fct, args) ->
+            Builtin.pureBuiltinFunctions.Contains fct.Name && List.forall isInlinableExpr args
+        | FunCall(Op op, args) -> not (Builtin.assignOps.Contains op) && List.forall isInlinableExpr args
+        | _ -> false
+
+    type [<NoComparison>] Inlining = {
+        func: TopLevel
+        argIndex: int
+        varDecl: VarDecl
+        argExpr: Expr
+    }
+
+    // Find when functions are always called with the same trivial expr, that can be inlined into the function body.
+    let findInlinings code: Inlining list =
+        let mutable argInlinings = []
+        Analyzer.resolve code
+        Analyzer.markWrites code
+        let funcInfos = Analyzer.findFuncInfos code
+        for funcInfo in funcInfos do
+            let canBeRenamed = not (options.noRenamingList |> List.contains funcInfo.name) // noRenamingList includes "main"
+            // If the function is overloaded, removing a parameter could conflict with another overload.
+            if canBeRenamed && not funcInfo.funcType.isExternal && funcInfo.isOverloaded then
+                let callSites = funcInfos |> List.collect (fun n -> n.callSites) |> List.filter (fun n -> n.prototype = funcInfo.funcType.prototype)
+                for argIndex, (_, argDecl) in List.indexed funcInfo.funcType.parameters do
+                    match argDecl.name.VarDecl with
+                    | Some varDecl when not varDecl.ty.isOutOrInout -> // Only inline 'in' parameters.
+                        let argExprs = callSites |> List.map (fun c -> c.argExprs |> List.item argIndex) |> List.distinct
+                        match argExprs with
+                        | [argExpr] when isInlinableExpr argExpr -> // The argExpr must always be the same at all call sites.
+                            argInlinings <- {func=funcInfo.func; argIndex=argIndex; varDecl=varDecl; argExpr=argExpr} :: argInlinings
+                        | _ -> ()
+                    | _ -> ()
+        argInlinings
+
+    let apply (didInline: bool ref) code =
+        let argInlinings = findInlinings code
+
+        let removeInlined func list =
+            list
+            |> List.indexed
+            |> List.filter (fun (idx, _) -> not (argInlinings |> List.exists (fun inl -> inl.func = func && inl.argIndex = idx)))
+            |> List.map snd
+
+        let applyExpr _ = function
+            | FunCall (Var v, argExprs) as f ->
+                // Remove the argument expression from the call site.
+                match v.Declaration with
+                | Declaration.Func fd -> FunCall (Var v, removeInlined fd.func argExprs)
+                | _ -> f
+            | x -> x
+
+        let applyTopLevel = function
+            | Function(fct, body) as f ->
+                // Handle argument inlining for other functions called by f.
+                let _, body = mapStmt (mapEnv applyExpr id) body
+                // Handle argument inlining for f. Remove the parameter from the declaration.
+                let fct = {fct with args = removeInlined f fct.args}
+                // Handle argument inlining for f. Insert in front of the body a declaration for each inlined argument.
+                let decls =
+                    argInlinings
+                    |> List.filter (fun inl -> inl.func = f)
+                    |> List.map (fun inl -> Decl (
+                        {inl.varDecl.ty with typeQ = inl.varDecl.ty.typeQ |> List.filter ((=) "const")},
+                        [{inl.varDecl.decl with init = Some inl.argExpr}]))
+                Function(fct, Block (decls @ body.asStmtList))
+            | tl -> tl
+
+        if argInlinings.IsEmpty then
+            code
+        else
+            let code = code |> List.map applyTopLevel
+            didInline.Value <- true
+            code
+
 let rec iterateSimplifyAndInline li =
     let li = if not options.noRemoveUnused then removeUnusedFunctions li else li
     if not options.noInlining then
@@ -600,14 +679,16 @@ let rec iterateSimplifyAndInline li =
         Analyzer.markInlinableFunctions li
         Analyzer.markInlinableVariables li
     let didInline = ref false
-    let simplified = mapTopLevel (mapEnv (simplifyExpr didInline) simplifyStmt) li
+    let li = mapTopLevel (mapEnv (simplifyExpr didInline) simplifyStmt) li
 
     // now that the functions were inlined, we can remove them
-    let simplified = simplified |> List.filter (function
+    let li = li |> List.filter (function
         | Function (funcType, _) -> not funcType.fName.ToBeInlined || funcType.fName.Name.StartsWith("i_")
         | _ -> true)
+
+    let li = if options.noInlining then li else ArgumentInlining.apply didInline li
 
-    if didInline.Value then iterateSimplifyAndInline simplified else simplified
+    if didInline.Value then iterateSimplifyAndInline li else li
 
 let simplify li =
     li

tests/commands.txt

@@ -9,7 +9,7 @@
 
 # Unit tests
 
---no-remove-unused --no-renaming --format c-variables -o tests/unit/blocks.expected tests/unit/blocks.frag
+--no-remove-unused --no-inlining --no-renaming --format c-variables -o tests/unit/blocks.expected tests/unit/blocks.frag
 --no-remove-unused --hlsl --no-inlining --no-renaming --format c-variables -o tests/unit/geometry.hlsl.expected tests/unit/geometry.hlsl
 --no-remove-unused --no-renaming --no-inlining --move-declarations --format c-array -o tests/unit/operators.expected tests/unit/operators.frag
 --no-renaming --no-inlining --format c-array -o tests/unit/minus-zero.expected tests/unit/minus-zero.frag
@@ -45,6 +45,7 @@
 --no-remove-unused --no-renaming --format indented --aggressive-inlining -o tests/unit/inline-fn.aggro.expected tests/unit/inline-fn.frag
 --no-remove-unused --format indented -o tests/unit/inline-fn-multiple.expected tests/unit/inline-fn-multiple.frag
 --no-remove-unused --no-renaming --format indented -o tests/unit/simplify.expected tests/unit/simplify.frag
+--no-remove-unused --no-renaming --format indented -o tests/unit/arg-inlining.expected tests/unit/arg-inlining.frag
 
 # Partial renaming tests
 

tests/compression_results.log

@@ -1,24 +1,24 @@
-clod.frag...                              8910 => 1525.318
-mouton/mouton.vert...                    17052 => 2453.185
-audio-flight-v2.frag                      4576 =>  898.648
-buoy.frag                                 4129 =>  626.967
-controllable-machinery.frag               7760 => 1219.338
-ed-209.frag                               7860 => 1355.308
+clod.frag...                              8897 => 1528.145
+mouton/mouton.vert...                    17026 => 2457.862
+audio-flight-v2.frag                      4538 =>  891.404
+buoy.frag                                 4097 =>  624.156
+controllable-machinery.frag               7720 => 1220.140
+ed-209.frag                               7768 => 1340.423
 elevated.hlsl                             3405 =>  603.218
-endeavour.frag                            2618 =>  534.518
-from-the-seas-to-the-stars.frag          14304 => 2332.852
-frozen-wasteland.frag                     4593 =>  807.919
+endeavour.frag                            2605 =>  534.116
+from-the-seas-to-the-stars.frag          14292 => 2328.094
+frozen-wasteland.frag                     4583 =>  806.518
 kinder_painter.frag                       2867 =>  447.669
-leizex.frag                               2311 =>  510.372
-lunaquatic.frag                           5257 => 1049.740
-mandelbulb.frag                           2400 =>  547.050
+leizex.frag                               2298 =>  510.191
+lunaquatic.frag                           5247 => 1049.030
+mandelbulb.frag                           2370 =>  538.322
 ohanami.frag                              3256 =>  722.517
 orchard.frag                              5537 => 1022.773
 oscars_chair.frag                         4651 =>  986.364
-robin.frag                                6310 => 1053.772
+robin.frag                                6297 => 1053.367
 slisesix.frag                             4573 =>  931.855
 terrarium.frag                            3626 =>  746.585
-the_real_party_is_in_your_pocket.frag    12160 => 1809.537
+the_real_party_is_in_your_pocket.frag    12111 => 1794.687
 valley_ball.glsl                          4386 =>  888.496
-yx_long_way_from_home.frag                2975 =>  610.699
-Total: 135516 => 23684.701
+yx_long_way_from_home.frag                2947 =>  606.406
+Total: 135097 => 23632.338