From e34b5a64f002cf297f83dc0761c0bd8996b56695 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20H=C3=B6ck?= <hock@zhaw.ch>
Date: Sat, 6 May 2023 09:35:17 +0200
Subject: [PATCH] [ codegen ] get rid of trivial let statements (#2961)

---
 CHANGELOG.md                       |  2 ++
 src/Compiler/Opts/ConstantFold.idr | 13 ++++++++-----
 tests/Main.idr                     |  2 +-
 tests/chez/constfold2/Check.idr    | 19 +++++++++++++++++++
 tests/chez/constfold2/Fold.idr     | 27 +++++++++++++++++++++++++++
 tests/chez/constfold2/expected     |  3 +++
 tests/chez/constfold2/input        |  2 ++
 tests/chez/constfold2/run          |  5 +++++
 8 files changed, 67 insertions(+), 6 deletions(-)
 create mode 100644 tests/chez/constfold2/Check.idr
 create mode 100644 tests/chez/constfold2/Fold.idr
 create mode 100644 tests/chez/constfold2/expected
 create mode 100644 tests/chez/constfold2/input
 create mode 100644 tests/chez/constfold2/run

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a40b27da..de61a75ea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -83,6 +83,8 @@
   `TTImp.ProcessDef.genRunTime`. This allows us to track when incomplete `case`
   blocks get the runtime error added.
 
+* Constant folding of trivial let statements.
+
 ### Library changes
 
 #### Prelude
diff --git a/src/Compiler/Opts/ConstantFold.idr b/src/Compiler/Opts/ConstantFold.idr
index cbd8781c1..87a60f3d7 100644
--- a/src/Compiler/Opts/ConstantFold.idr
+++ b/src/Compiler/Opts/ConstantFold.idr
@@ -69,6 +69,11 @@ lookup fc (MkVar p) rho = case go p rho of
       Z => Left (MkVar First)
       S i' => bimap later weaken (go (dropLater p) (Wk ws' rho))
 
+replace : CExp vars -> Bool
+replace (CLocal _ _)   = True
+replace (CPrimVal _ _) = True
+replace _              = False
+
 -- constant folding of primitive functions
 -- if a primitive function is applied to only constant
 -- then replace with the result
@@ -83,11 +88,9 @@ constFold rho (CLam fc x y)
   = CLam fc x $ constFold (wk (mkSizeOf [x]) rho) y
 constFold rho (CLet fc x inl y z) =
     let val = constFold rho y
-     in case val of
-        CPrimVal _ _ => if inl == YesInline
-            then constFold (val :: rho) z
-            else CLet fc x inl val (constFold (wk (mkSizeOf [x]) rho) z)
-        _ => CLet fc x inl val (constFold (wk (mkSizeOf [x]) rho) z)
+     in if replace val
+          then constFold (val :: rho) z
+          else CLet fc x inl val (constFold (wk (mkSizeOf [x]) rho) z)
 constFold rho (CApp fc (CRef fc2 n) [x]) =
   if n == NS typesNS (UN $ Basic "prim__integerToNat")
      then case constFold rho x of
diff --git a/tests/Main.idr b/tests/Main.idr
index a7def42d5..b61b7aa7b 100644
--- a/tests/Main.idr
+++ b/tests/Main.idr
@@ -310,7 +310,7 @@ chezTests = MkTestPool "Chez backend" [] (Just Chez)
     , "futures001"
     , "bitops"
     , "casts"
-    , "constfold"
+    , "constfold", "constfold2"
     , "memo"
     , "newints"
     , "integers"
diff --git a/tests/chez/constfold2/Check.idr b/tests/chez/constfold2/Check.idr
new file mode 100644
index 000000000..8d6fcad97
--- /dev/null
+++ b/tests/chez/constfold2/Check.idr
@@ -0,0 +1,19 @@
+import Data.List
+import Data.String
+import System.File
+
+path : String
+path = "build/exec/fold_app/fold.ss"
+
+mainLine : String -> Bool
+mainLine str =
+  ("(define Main-main" `isPrefixOf` str) &&
+  not ("Main-manyToSmall" `isInfixOf` str)
+
+main : IO ()
+main = do
+  Right str <- readFile path
+    | Left err => putStrLn "Error when reading \{path}"
+  case any mainLine (lines str) of
+    True  => putStrLn "Constant expression correctly folded"
+    False => putStrLn "Failed to fold constant expression"
diff --git a/tests/chez/constfold2/Fold.idr b/tests/chez/constfold2/Fold.idr
new file mode 100644
index 000000000..88b30bcbe
--- /dev/null
+++ b/tests/chez/constfold2/Fold.idr
@@ -0,0 +1,27 @@
+import Data.Nat
+
+%default total
+
+record Small where
+  constructor S
+  value : Nat
+  {auto 0 prf : LTE value 20}
+
+Show Small where show = show . value
+
+record Smaller where
+  constructor XS
+  value : Nat
+  {auto 0 prf : LTE value 10}
+
+-- This is the identity function
+toSmall : Smaller -> Small
+toSmall (XS v @{prf}) = S v @{transitive prf %search}
+
+-- This is again the identity function
+manyToSmall : List Smaller -> List Small
+manyToSmall [] = []
+manyToSmall (x::xs) = toSmall x :: manyToSmall xs
+
+main : IO ()
+main = printLn $ manyToSmall [XS 1, XS 2, XS 3]
diff --git a/tests/chez/constfold2/expected b/tests/chez/constfold2/expected
new file mode 100644
index 000000000..049d27828
--- /dev/null
+++ b/tests/chez/constfold2/expected
@@ -0,0 +1,3 @@
+1/1: Building Check (Check.idr)
+Main> Constant expression correctly folded
+Main> Bye for now!
diff --git a/tests/chez/constfold2/input b/tests/chez/constfold2/input
new file mode 100644
index 000000000..fc5992c29
--- /dev/null
+++ b/tests/chez/constfold2/input
@@ -0,0 +1,2 @@
+:exec main
+:q
diff --git a/tests/chez/constfold2/run b/tests/chez/constfold2/run
new file mode 100644
index 000000000..302d8691d
--- /dev/null
+++ b/tests/chez/constfold2/run
@@ -0,0 +1,5 @@
+rm -rf build
+
+$1 --no-banner --no-color --quiet -o fold Fold.idr
+$1 --no-banner --no-color --console-width 0 Check.idr < input
+