Allow inlining one-line functions

This implements the simplest possible implementation of inlining
functions. It only supports a function consisting of a single `return`,
when it has been explicitly named as `i_whatever`. Basically, this works
in cases where you might want to use a macro, excpet this gives the
minifier full visibility through it.

The simplicity of this implementation means that it's likely possible to
get into trouble by marking certain functions as to-be-inlined. In
particular, I'm pretty sure the following will break it:
- inlining overloaded functions
- inlining a function into a scope which has a local shadowing a global,
  where the inlined function refers to that global (though this will
  also break a macro)

README.md

@@ -237,8 +237,8 @@ This is enabled with `--aggressive-inlining`.
 
 ### Explicit inlining
 
-Shader Minifier will always inline variables that starts with `i_`. Inlining can allow
-the Minifier to simplify the code further.
+Shader Minifier will always inline variables and functions that start with
+`i_`. Inlining can allow the Minifier to simplify the code further.
 
 For example, this input:
 
@@ -264,6 +264,31 @@ int foo(int x,int y)
 }
 ```
 
+And this input:
+
+```c
+float i_foo(float f, float g, float x) {
+  return f*x + g*x + f*g;
+}
+
+float bar(float a) {
+  return i_foo(2.0, 3.0, sin(sqrt(a)));
+}
+```
+
+will be simplified into:
+
+```c
+float bar(float a)
+{
+  return 2.*sin(sqrt(a))+3.*sin(sqrt(a))+6.;
+}
+```
+
+Note that the function inlining is very simplistic, only supporting functions
+which consist of a single return statement. (Basically, cases where you would
+use a macro, except this gives the minifier full visibility through it.)
+
 If you want to aggressively reduce the size of your shader, try inlining more
 variables. Inlining can have performance implications though (if the variable
 stored the result of a computation), so be careful with it.

src/ast.fs

@@ -131,7 +131,7 @@ type MapEnv = {
     fExpr: MapEnv -> Expr -> Expr
     fStmt: Stmt -> Stmt
     vars: Map<string, Type * DeclElt>
-    fns: Map<string, FunctionType>
+    fns: Map<string, FunctionType * Stmt>
 }
 
 let mapEnv fe fi = {fExpr = fe; fStmt = fi; vars = Map.empty; fns = Map.empty}
@@ -211,7 +211,7 @@ let mapTopLevel env li =
             let env, res = mapDecl env t
             env, TLDecl res
         | Function(fct, body) ->
-            let env = {env with fns = env.fns.Add(fct.fName.Name, fct)}
+            let env = {env with fns = env.fns.Add(fct.fName.Name, (fct, body))}
             env, Function(fct, snd (mapStmt env body))
         | e -> env, e)
     res

src/rewriter.fs

@@ -63,7 +63,34 @@ let private bool = function
     | true -> Var (Ident "true") // Int (1, "")
     | false -> Var (Ident "false") // Int (0, "")
 
+let private inlineFn (declArgs:Decl list) passedArgs bodyExpr =
+    let mutable argMap = Map.empty
+    for declArg, passedArg in List.zip declArgs passedArgs do
+        let declElt = List.exactlyOne (snd declArg)
+        argMap <- argMap.Add(declElt.name.Name, passedArg)
+    let mapInline _ = function
+        | Var iv as ie ->
+            match argMap.TryFind iv.Name with
+            | Some inlinedExpr -> inlinedExpr
+            | _ -> ie
+        | ie -> ie
+    mapExpr (mapEnv mapInline id) bodyExpr
+
 let rec private simplifyExpr didInline env = function
+    | FunCall(Var v, passedArgs) as e when v.ToBeInlined ->
+        match env.fns.TryFind v.Name with
+        | None -> e
+        | Some ({args = declArgs}, body) ->
+            match body with
+            | Jump (JumpKeyword.Return, Some bodyExpr)
+            | Block [Jump (JumpKeyword.Return, Some bodyExpr)] ->
+                didInline := true
+                inlineFn declArgs passedArgs bodyExpr
+            // Don't yell if we've done some inlining this pass -- maybe it
+            // turned the function into a one-liner, so allow trying again on
+            // the next pass. (If it didn't, we'll yell next pass.)
+            | _ when !didInline -> e
+            | _ -> failwithf "Cannot inline %s since it consists of more than a single return" v.Name
     | FunCall(Op "-", [Int (i1, su)]) -> Int (-i1, su)
     | FunCall(Op "-", [FunCall(Op "-", [e])]) -> e
     | FunCall(Op "+", [e]) -> e
@@ -329,7 +356,7 @@ let markLValues li =
             FunCall(Op o, (mapExpr newEnv e)::args)
         | FunCall(Var v, _) as e ->
             match env.fns.TryFind v.Name with
-            | Some fct when fct.fName.IsLValue ->
+            | Some (fct, _) when fct.fName.IsLValue ->
                 let newEnv = {env with fExpr = markVars}
                 mapExpr newEnv e
             | _ -> e
@@ -360,11 +387,12 @@ let simplify li =
     |> if options.aggroInlining then markLValues >> inlineAllConsts else id
     |> reorderTopLevel
     |> iterateSimplifyAndInline
-    |> List.map (function
-        | TLDecl (ty, li) -> TLDecl (rwType ty, declsNotToInline li)
-        | TLVerbatim s -> TLVerbatim (stripSpaces s)
-        | Function (fct, body) -> Function (rwFType fct, body)
-        | e -> e
+    |> List.choose (function
+        | TLDecl (ty, li) -> TLDecl (rwType ty, declsNotToInline li) |> Some
+        | TLVerbatim s -> TLVerbatim (stripSpaces s) |> Some
+        | Function (fct, _) when fct.fName.ToBeInlined -> None
+        | Function (fct, body) -> Function (rwFType fct, body) |> Some
+        | e -> e |> Some
     )
     |> squeezeTLDeclarations
 

tests/commands.txt

@@ -12,12 +12,17 @@
 --hlsl --no-inlining --no-renaming --format c-variables -o tests/unit/geometry.hlsl.expected tests/unit/geometry.hlsl
 --no-renaming --no-inlining --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
+--no-renaming --format c-array --no-inlining -o tests/unit/float.frag.expected tests/unit/float.frag
+
+# Inlining unit tests
+
 --no-renaming --format indented -o tests/unit/inline.expected tests/unit/inline.frag
 --no-renaming --format indented --aggressive-inlining -o tests/unit/inline.aggro.expected tests/unit/inline.frag
 --no-renaming --format indented --no-inlining -o tests/unit/inline.no.expected tests/unit/inline.frag
 --no-renaming --format indented -o tests/unit/inline-aggro.expected tests/unit/inline-aggro.frag
 --no-renaming --format indented --aggressive-inlining -o tests/unit/inline-aggro.aggro.expected tests/unit/inline-aggro.frag
---no-renaming --format c-array --no-inlining -o tests/unit/float.frag.expected tests/unit/float.frag
+--no-renaming --format indented -o tests/unit/inline-fn.expected tests/unit/inline-fn.frag
+--no-renaming --format indented --aggressive-inlining -o tests/unit/inline-fn.aggro.expected tests/unit/inline-fn.frag
 
 # Partial renaming tests
 

tests/unit/inline-fn.aggro.expected

@@ -0,0 +1,21 @@
+float a()
+{
+  return vec3(2.,3.,4.).x+vec3(2.,3.,4.).y;
+}
+float b(float g)
+{
+  float x=(g+10.)*vec3(20.,30.,40.).x+vec3(20.,30.,40.).y,y=(g+10.1)*vec3(20.1,30.1,40.1).x+vec3(20.1,30.1,40.1).y;
+  return x+y;
+}
+float c()
+{
+  return 154.;
+}
+float d()
+{
+  return 54.;
+}
+float e()
+{
+  return 84.;
+}

tests/unit/inline-fn.expected

@@ -0,0 +1,22 @@
+float a()
+{
+  return vec3(2.,3.,4.).x+vec3(2.,3.,4.).y;
+}
+float b(float g)
+{
+  float x=(g+10.)*vec3(20.,30.,40.).x+vec3(20.,30.,40.).y,y=(g+10.1)*vec3(20.1,30.1,40.1).x+vec3(20.1,30.1,40.1).y;
+  return x+y;
+}
+float c()
+{
+  return 154.;
+}
+float d()
+{
+  float f=7.,g=f*f+1.;
+  return g+4.;
+}
+float e()
+{
+  return 84.;
+}

tests/unit/inline-fn.frag

@@ -0,0 +1,37 @@
+float i_foo(float f, vec3 g) {
+  return f*g.x + g.y;
+}
+
+float a() {
+  return i_foo(1.0, vec3(2.0, 3.0, 4.0));
+}
+
+float b(float g) {
+  float x = i_foo(g + 10.0, vec3(20.0, 30.0, 40.0));
+  float y = i_foo(g + 10.1, vec3(20.1, 30.1, 40.1));
+  return x + y;
+}
+
+float i_bar(float f, float g) {
+  return f*g + 1.0;
+}
+
+float c() {
+  return i_bar(i_bar(2.0, 3.0), i_bar(4.0, 5.0)) + 6.0;
+}
+
+float d() {
+  float f = i_bar(2.0, 3.0);
+  float g = i_bar(f, f);
+  return g + 4.0;
+}
+
+float i_multipass(float x) {
+  float f = 42.0;
+  float g = 2.0;
+  return f*g*x;
+}
+
+float e() {
+  return i_multipass(1.0);
+}