Reuse function parameters inside the function

src/ast.fs

@@ -191,7 +191,7 @@ type Shader = {
 // while also collecting visible variable and function declarations along the way.
 
 [<NoComparison>] [<RequireQualifiedAccess>]
-type BlockLevel = FunctionRoot | Nested | Unknown
+type BlockLevel = FunctionRoot of FunctionType | Nested | Unknown
 
 [<NoComparison; NoEquality>]
 type MapEnv = {
@@ -315,7 +315,7 @@ let mapTopLevel env li =
             let env = env.withFunction(fct, body, replaceMostRecentOverload = true)
 
             // Transform the body. The env modifications (local variables) are discarded.
-            let _, body = mapStmt BlockLevel.FunctionRoot env body
+            let _, body = mapStmt (BlockLevel.FunctionRoot fct) env body
             // Update env.fns with the transformed body.
             let env = env.withFunction(fct, body, replaceMostRecentOverload = true)
 

src/rewriter.fs

@@ -428,7 +428,7 @@ module private RewriterImpl =
     // Reuse an existing local variable declaration that won't be used anymore, instead of introducing a new one.
     // The reused identifier gets compressed better, and the declaration is sometimes removed.
     // float d1 = f(); float d2 = g();  ->  float d1 = f(); d1 = g();
-    let reuseExistingVarDecl b =
+    let reuseExistingVarDecl blockLevel b =
         let tryReplaceWithPrecedingAndFollowing f (xs : Stmt list) =
             let rec go f preceding = function
                 | head :: tail ->
@@ -440,20 +440,22 @@ module private RewriterImpl =
         b |> tryReplaceWithPrecedingAndFollowing (function
         | (preceding2, Decl (ty2, declElts), following2) ->
             let findAssignmentReplacementFor declElt2 declBefore2 declAfter2 =
-                let compatibleDeclElts = (preceding2 @ declBefore2) |> List.collect (function
-                    // The previous declaration must have the same type.
-                    | Decl (ty1, declElts1) when ty2 = ty1 ->
-                        let firstIsNotUsedAfterSecondDeclared (declElt1 : DeclElt) followingDecl2 =
-                            // The first variable must not be used after the second is declared.
-                            Analyzer.varUsesInStmt (Block followingDecl2) |> List.forall (fun i -> i.Name <> declElt1.name.Name)
-                        declElts1 |> List.filter (fun declElt1 ->
-                            // The previous declaration must have the same size and semantics
-                            declElt1.size = declElt2.size &&
-                            declElt1.semantics = declElt2.semantics &&
-                            firstIsNotUsedAfterSecondDeclared declElt1 (declAfter2 @ following2)
-                        )
-                    | _ -> [])
-                match compatibleDeclElts |> List.tryHead with
+                // Collect previous declarations of the same type.
+                let localDecls = (preceding2 @ declBefore2) |> List.collect (function Decl (ty1, declElts1) when ty2 = ty1 -> declElts1 | _ -> [])
+                let args =
+                    match blockLevel with
+                    | BlockLevel.FunctionRoot fct ->
+                        fct.parameters |> List.choose (function ty, decl -> if not ty.isOutOrInout && ty = ty2 then Some decl else None)
+                    | _ -> []
+
+                let compatibleDeclElt = (localDecls @ args) |> List.tryFind (fun declElt1 ->
+                    declElt1.size = declElt2.size &&
+                    declElt1.semantics = declElt2.semantics &&
+                    // The first variable must not be used after the second is declared.
+                    Analyzer.varUsesInStmt (Block (declAfter2 @ following2)) |> List.forall (fun i -> i.Name <> declElt1.name.Name)
+                )
+
+                match compatibleDeclElt with
                 | None -> None
                 | Some declElt1 ->
                     debug $"{declElt2.name.Loc}: eliminating local variable '{declElt2.name}' by reusing existing local variable '{declElt1.name}'"
@@ -592,9 +594,9 @@ module private RewriterImpl =
 
         let b =
             // We ensure control flow analysis is trivial by only doing this at the root block level.
-            if blockLevel <> BlockLevel.FunctionRoot || hasPreprocessor
+            if (match blockLevel with BlockLevel.FunctionRoot _ -> false | _ -> true) || hasPreprocessor
             then b
-            else reuseExistingVarDecl b
+            else reuseExistingVarDecl blockLevel b
 
         // Consecutive declarations of the same type become one.  float a;float b;  ->  float a,b;
         let b = squeezeConsecutiveDeclarations b
@@ -772,7 +774,7 @@ module private ArgumentInlining =
         let applyTopLevel = function
             | Function(fct, body) as f ->
                 // Handle argument inlining for other functions called by f.
-                let _, body = mapStmt BlockLevel.FunctionRoot (mapEnvExpr applyExpr) body
+                let _, body = mapStmt (BlockLevel.FunctionRoot fct) (mapEnvExpr applyExpr) 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.

tests/compression_results.log

@@ -1,24 +1,24 @@
-clod.frag...                              8842 => 1513.669
-mouton/mouton.vert...                    16972 => 2439.859
-audio-flight-v2.frag                      4526 =>  883.304
-buoy.frag                                 4070 =>  619.904
-controllable-machinery.frag               7708 => 1220.329
-ed-209.frag                               7734 => 1333.865
+clod.frag...                              8836 => 1515.830
+mouton/mouton.vert...                    16951 => 2447.462
+audio-flight-v2.frag                      4509 =>  881.311
+buoy.frag                                 4051 =>  607.428
+controllable-machinery.frag               7703 => 1216.907
+ed-209.frag                               7702 => 1341.619
 elevated.hlsl                             3405 =>  603.218
 endeavour.frag                            2589 =>  529.811
-from-the-seas-to-the-stars.frag          14250 => 2312.836
+from-the-seas-to-the-stars.frag          14250 => 2308.842
 frozen-wasteland.frag                     4566 =>  806.475
 kinder_painter.frag                       2847 =>  442.771
-leizex.frag                               2275 =>  509.430
-lunaquatic.frag                           5235 => 1044.000
+leizex.frag                               2270 =>  510.375
+lunaquatic.frag                           5239 => 1044.087
 mandelbulb.frag                           2347 =>  537.664
-ohanami.frag                              3255 =>  723.026
-orchard.frag                              5536 => 1022.550
+ohanami.frag                              3255 =>  723.732
+orchard.frag                              5510 => 1027.042
 oscars_chair.frag                         4651 =>  986.364
-robin.frag                                6268 => 1043.078
+robin.frag                                6257 => 1047.684
 slisesix.frag                             4549 =>  913.642
-terrarium.frag                            3611 =>  744.367
-the_real_party_is_in_your_pocket.frag    12101 => 1787.736
-valley_ball.glsl                          4334 =>  890.606
-yx_long_way_from_home.frag                2926 =>  599.317
-Total: 134597 => 23507.820
+terrarium.frag                            3593 =>  741.557
+the_real_party_is_in_your_pocket.frag    12087 => 1780.714
+valley_ball.glsl                          4307 =>  881.820
+yx_long_way_from_home.frag                2942 =>  598.650
+Total: 134416 => 23495.003

tests/real/audio-flight-v2.frag.expected

@@ -51,13 +51,13 @@ float vmax(vec3 v)
 }
 float fBox(vec3 p,vec3 b)
 {
-  vec3 d=abs(p)-b;
-  return length(max(d,vec3(0)))+vmax(min(d,vec3(0)));
+  p=abs(p)-b;
+  return length(max(p,vec3(0)))+vmax(min(p,vec3(0)));
 }
 float fBox2(vec2 p,vec2 b)
 {
-  vec2 d=abs(p)-b;
-  return length(max(d,vec2(0)))+vmax(min(d,vec2(0)));
+  p=abs(p)-b;
+  return length(max(p,vec2(0)))+vmax(min(p,vec2(0)));
 }
 float sdCap(vec3 p)
 {
@@ -188,9 +188,9 @@ vec2 marcher(vec3 ro,vec3 rd)
 }
 vec3 normal(vec3 p,float t)
 {
-  float e=MINDIST*t;
+  t*=MINDIST;
   vec2 h=vec2(1,-1)*.5773;
-  return normalize(h.xyy*map(p+h.xyy*e,0.).x+h.yyx*map(p+h.yyx*e,0.).x+h.yxy*map(p+h.yxy*e,0.).x+h.xxx*map(p+h.xxx*e,0.).x);
+  return normalize(h.xyy*map(p+h.xyy*t,0.).x+h.yyx*map(p+h.yyx*t,0.).x+h.yxy*map(p+h.yxy*t,0.).x+h.xxx*map(p+h.xxx*t,0.).x);
 }
 vec3 hsv2rgb(vec3 c)
 {

tests/real/buoy.frag.expected

@@ -77,17 +77,16 @@ float WavesSmooth(vec3 pos)
 }
 float WaveCrests(vec3 ipos,vec2 fragCoord)
 {
-  vec3 pos=ipos;
-  pos*=.2*vec3(1);
+  ipos*=.2*vec3(1);
   float f=0.;
-  pos+=iTime*vec3(0,.1,.1);
-  vec3 pos2=pos;
+  ipos+=iTime*vec3(0,.1,.1);
+  vec3 pos2=ipos;
   for(int i=0;i<6;i++)
-    pos=(pos.yzx+pos.zyx*vec3(1,-1,1))/sqrt(2.),f=f*1.5+abs(Noise(pos).x-.5)*2.,pos*=2.;
-  pos=pos2*exp2(float(6));
-  pos.y=-.05*iTime;
+    ipos=(ipos.yzx+ipos.zyx*vec3(1,-1,1))/sqrt(2.),f=f*1.5+abs(Noise(ipos).x-.5)*2.,ipos*=2.;
+  ipos=pos2*exp2(float(6));
+  ipos.y=-.05*iTime;
   for(int i=6;i<16;i++)
-    pos=(pos.yzx+pos.zyx*vec3(1,-1,1))/sqrt(2.),f=f*1.5+pow(abs(Noise(pos).x-.5)*2.,1.),pos*=2.;
+    ipos=(ipos.yzx+ipos.zyx*vec3(1,-1,1))/sqrt(2.),f=f*1.5+pow(abs(Noise(ipos).x-.5)*2.,1.),ipos*=2.;
   f/=1500.;
   f-=Noise(ivec2(fragCoord.xy)).x*.01;
   return pow(smoothstep(.4,-.1,f),6.);
@@ -145,8 +144,8 @@ vec3 ShadeBoat(vec3 pos,vec3 ray)
   albedo=mix(vec3(1,.8,.08),albedo,smoothstep(.05-aa,.05,abs(abs(pos.y)-.65)));
   albedo*=smoothstep(-.1,1.,dot(norm,lightDir))*vec3(1,.9,.8)+vec3(.06,.1,.1);
   lightDir=pow(max(0.,dot(norm,normalize(lightDir-ray))),1e2)*1e2/32.*vec3(1);
-  vec3 rr=reflect(ray,norm);
-  lightDir+=mix(vec3(0,.04,.04),Sky(rr),smoothstep(-.1,.1,rr.y));
+  pos=reflect(ray,norm);
+  lightDir+=mix(vec3(0,.04,.04),Sky(pos),smoothstep(-.1,.1,pos.y));
   aa=pow(1.-abs(dot(norm,ray)),5.);
   aa=mix(.001,1.,aa);
   return mix(albedo,lightDir,aa);
@@ -189,15 +188,15 @@ vec3 ShadeOcean(vec3 pos,vec3 ray,vec2 fragCoord)
   vec3 reflectedRay=ray-2.*norm*ndotr;
   norm=ray+(-cos(1.33*acos(-ndotr))-ndotr)*norm;
   norm=normalize(norm);
-  vec3 reflection=Sky(reflectedRay);
+  ray=Sky(reflectedRay);
   float t=TraceBoat(pos-0.*reflectedRay,reflectedRay);
   if(t>0.)
-    reflection=ShadeBoat(pos+t*reflectedRay,reflectedRay);
+    ray=ShadeBoat(pos+t*reflectedRay,reflectedRay);
   t=TraceBoat(pos-0.*norm,norm);
   reflectedRay=vec3(0,.04,.04);
   if(t>0.)
     reflectedRay=mix(reflectedRay,ShadeBoat(pos+t*norm,norm),exp(-t));
-  reflectedRay=mix(reflectedRay,reflection,pow(1.-abs(ndotr),5.));
+  reflectedRay=mix(reflectedRay,ray,pow(1.-abs(ndotr),5.));
   return mix(reflectedRay,vec3(1),WaveCrests(pos,fragCoord));
 }
 void mainImage(out vec4 fragColor,vec2 fragCoord)

tests/real/controllable-machinery.frag.expected

@@ -40,8 +40,8 @@ vec4 BPos(float t)
 }
 float PrBoxDf(vec3 p,vec3 b)
 {
-  vec3 d=abs(p)-b;
-  return min(max(d.x,max(d.y,d.z)),0.)+length(max(d,0.));
+  p=abs(p)-b;
+  return min(max(p.x,max(p.y,p.z)),0.)+length(max(p,0.));
 }
 float PrRoundBoxDf(vec3 p,vec3 b)
 {

tests/real/ed-209.frag.expected

@@ -10,17 +10,18 @@ float stretch,gunsUp,gunsForward,edWalk,edTwist,edDown,edShoot,doorOpen,glow;
 struct MarchData{float d;vec3 mat;float specPower;};
 mat2 rot(float a)
 {
-  float c=cos(a),s=sin(a);
-  return mat2(c,s,-s,c);
+  float c=cos(a);
+  a=sin(a);
+  return mat2(c,a,-a,c);
 }
 float remap(float f,float in1,float in2,float out1,float out2)
 {
   return mix(out1,out2,clamp((f-in1)/(in2-in1),0.,1.));
 }
 float sdBox(vec3 p,vec3 b)
 {
-  vec3 q=abs(p)-b;
-  return length(max(q,0.))+min(max(q.x,max(q.y,q.z)),0.);
+  p=abs(p)-b;
+  return length(max(p,0.))+min(max(p.x,max(p.y,p.z)),0.);
 }
 float sdChamferedCube(vec3 p,vec3 r)
 {
@@ -54,8 +55,9 @@ float sdCappedCylinder(vec3 p,float h,float r)
 }
 float sdCapsule(vec3 p,vec3 a,vec3 b)
 {
-  vec3 pa=p-a,ba=b-a;
-  return length(pa-ba*clamp(dot(pa,ba)/dot(ba,ba),0.,1.))-.2;
+  p-=a;
+  a=b-a;
+  return length(p-a*clamp(dot(p,a)/dot(a,a),0.,1.))-.2;
 }
 float sdOctogon(vec2 p,float r)
 {
@@ -68,8 +70,9 @@ float sdOctogon(vec2 p,float r)
 }
 vec3 getRayDir(vec3 ro,vec3 lookAt,vec2 uv)
 {
-  vec3 forward=normalize(lookAt-ro),right=normalize(cross(vec3(0,1,0),forward));
-  return normalize(forward+right*uv.x+cross(forward,right)*uv.y);
+  ro=normalize(lookAt-ro);
+  lookAt=normalize(cross(vec3(0,1,0),ro));
+  return normalize(ro+lookAt*uv.x+cross(ro,lookAt)*uv.y);
 }
 MarchData minResult(MarchData a,MarchData b)
 {
@@ -306,11 +309,11 @@ MarchData map(vec3 p)
 }
 float calcShadow(vec3 p,vec3 lightPos)
 {
-  vec3 rd=normalize(lightPos-p);
+  lightPos=normalize(lightPos-p);
   float res=1.,t=.1;
   for(float i=0.;i<SHADOW_STEPS;i++)
     {
-      float h=map(p+rd*t).d;
+      float h=map(p+lightPos*t).d;
       res=min(res,12.*h/t);
       t+=h;
       if(res<.001||t>25.)
@@ -325,17 +328,17 @@ vec3 calcNormal(vec3 p,float t)
 }
 vec3 vignette(vec3 col,vec2 fragCoord)
 {
-  vec2 q=fragCoord.xy/iResolution.xy;
-  col*=.5+.5*pow(16.*q.x*q.y*(1.-q.x)*(1.-q.y),.4);
+  fragCoord=fragCoord.xy/iResolution.xy;
+  col*=.5+.5*pow(16.*fragCoord.x*fragCoord.y*(1.-fragCoord.x)*(1.-fragCoord.y),.4);
   return col;
 }
 vec3 applyLighting(vec3 p,vec3 rd,float d,MarchData data)
 {
   const vec3 sunPos=vec3(10,10,-10);
   vec3 sunDir=normalize(sunPos-p),n=calcNormal(p,d);
-  float primary=max(0.,dot(sunDir,n));
-  primary*=mix(.2,1.,calcShadow(p,sunPos));
-  sunDir=((primary+max(0.,dot(-sunDir,n))*.2)*(map(p+.33*n).d/.33)+pow(max(0.,dot(rd,reflect(sunDir,n))),data.specPower)*2.)*vec3(2,1.6,1.7);
+  d=max(0.,dot(sunDir,n));
+  d*=mix(.2,1.,calcShadow(p,sunPos));
+  sunDir=((d+max(0.,dot(-sunDir,n))*.2)*(map(p+.33*n).d/.33)+pow(max(0.,dot(rd,reflect(sunDir,n))),data.specPower)*2.)*vec3(2,1.6,1.7);
   return mix(data.mat*sunDir,vec3(.01),smoothstep(.7,1.,1.+dot(rd,n)))*exp(-length(p)*.05);
 }
 vec3 getSceneColor(vec3 ro,vec3 rd)

tests/real/from-the-seas-to-the-stars.frag.expected

@@ -12,8 +12,9 @@ float s(vec2 p)
 }
 float valnoise(vec2 p)
 {
-  vec2 c=floor(p),f=smoothstep(0.,1.,fract(p));
-  return mix(mix(s(c+vec2(0)),s(c+vec2(1,0)),f.x),mix(s(c+vec2(0,1)),s(c+vec2(1)),f.x),f.y);
+  vec2 c=floor(p);
+  p=smoothstep(0.,1.,fract(p));
+  return mix(mix(s(c+vec2(0)),s(c+vec2(1,0)),p.x),mix(s(c+vec2(0,1)),s(c+vec2(1)),p.x),p.y);
 }
 float fbm(vec2 p)
 {
@@ -49,18 +50,19 @@ float koi(vec2 t)
     return 1.;
   if(abs(t.y)>.5)
     return 1.;
-  vec2 p=t.xy*1.8,oop=p;
-  p.y*=.8;
-  p.x*=.9;
-  p.y+=pow(clamp(p.x-.7,0,1),2)*.2;
-  vec2 op=p;
-  if(p.y<0)
-    p.y*=mix(1.5,1.6,clamp(p.x,0,1));
-  p.y*=mix(1.6,1.,.5+.5*cos(clamp((p.x-.6)*11,0.,6.2831852)));
-  p.y-=p.x*.02;
-  p.x=mix(p.x,-.4,pow(clamp((p.x-.6)*1.1,0,1),1.4));
-  float d=abs(p.y)+p.x*p.x*.35-.34;
-  p.x-=.9;
+  t=t.xy*1.8;
+  vec2 oop=t;
+  t.y*=.8;
+  t.x*=.9;
+  t.y+=pow(clamp(t.x-.7,0,1),2)*.2;
+  vec2 op=t;
+  if(t.y<0)
+    t.y*=mix(1.5,1.6,clamp(t.x,0,1));
+  t.y*=mix(1.6,1.,.5+.5*cos(clamp((t.x-.6)*11,0.,6.2831852)));
+  t.y-=t.x*.02;
+  t.x=mix(t.x,-.4,pow(clamp((t.x-.6)*1.1,0,1),1.4));
+  float d=abs(t.y)+t.x*t.x*.35-.34;
+  t.x-=.9;
   op.y-=.05;
   op.y=abs(op.y);
   d=max(d,op.x-pow(op.y/2,.7)-1.1);