Tree shaker

app/Main.hs

@@ -22,6 +22,7 @@ import Data.Void qualified
 import GHC.IO.Handle (hFlush)
 import GHC.IO.StdHandles (stdout)
 import GitHash
+import Optimizer
 import Options.Applicative
 import Parser
 import System.Console.Repline
@@ -144,7 +145,8 @@ main = do
                     Right expr -> return expr
 
                 when debug $ putStrLn $ prettyPrintProgram expr
-                potentiallyTimedOperation "Compilation" showTime (evalStateT (BytecodeCompiler.compileProgram expr) (BytecodeCompiler.initCompilerState expr))
+                potentiallyTimedOperation "Compilation" showTime $ do
+                    optimize <$> evalStateT (BytecodeCompiler.compileProgram expr) (BytecodeCompiler.initCompilerState expr)
             else do
                 bytecode <- inputFileBinary input
                 return $ VM.fromBytecode bytecode

indigo.cabal

@@ -31,6 +31,7 @@ library
                      , Util
                      , RegInst
                      , Ffi
+                     , Optimizer
     other-modules: Paths_indigo
     build-depends: base,
         containers,
@@ -104,6 +105,7 @@ test-suite spec
         IntegrationSpec
         VerifierSpec
         RegInstSpec
+        OptimizerSpec
     build-tool-depends: hspec-discover:hspec-discover == 2.*
     default-language: Haskell2010
     default-extensions: OverloadedStrings, ImportQualifiedPost, QuasiQuotes, DisambiguateRecordFields, DuplicateRecordFields, OverloadedRecordDot

lib/BytecodeCompiler.hs

@@ -82,7 +82,7 @@ compileProgram (Parser.Program expr) = do
         i <- liftIO preludeFile
         case parseProgram (T.pack i) CompilerFlags{verboseMode = False, needsMain = False} of -- FIXME: pass on flags
             Left err -> error $ "Parse error: " ++ errorBundlePretty err
-            Right (Parser.Program progExpr) -> return progExpr
+            Right (Parser.Program progExpr) -> return $ progExpr ++ [Parser.FuncDef "__sep" [] Parser.Placeholder]
     prelude' <- concatMapM compileExpr prelude
     freePart <- concatMapM compileExpr expr
     createVirtualFunctions
@@ -407,11 +407,13 @@ compileExpr (Parser.Import{objects = o, from = from, as = as, qualified = qualif
     let expr = case parseProgram (T.pack i) Parser.initCompilerFlags of -- FIXME: pass on flags
             Left err -> error $ "Parse error: " ++ errorBundlePretty err
             Right (Parser.Program exprs) -> exprs
-    if qualified || isJust as
-        then do
-            let alias = if qualified then from else fromJust as
-            concatMapM compileExpr (map (`mangleAST` alias) expr)
-        else concatMapM compileExpr expr
+    let p =
+            if qualified || isJust as
+                then do
+                    let alias = if qualified then from else fromJust as
+                    concatMapM compileExpr (map (`mangleAST` alias) expr)
+                else concatMapM compileExpr expr
+    p >>= \p' -> return $ p' ++ [Label "__sep"]
   where
     mangleAST :: Parser.Expr -> String -> Parser.Expr
     mangleAST (Parser.FuncDec name types) alias = Parser.FuncDec (alias ++ "@" ++ name) types

lib/Optimizer.hs

@@ -0,0 +1,74 @@
+module Optimizer where
+
+import VM
+
+treeShake :: [Instruction] -> [Instruction]
+treeShake prog = do
+    let grouped = groupByLabel prog
+    let mainProg = dropWhile (\case (Label l : _) -> labelBaseName l /= "__sep"; _ -> False) grouped
+    let usedByDeps = shake mainProg grouped
+    concat usedByDeps
+  where
+    shakeOnce :: [[Instruction]] -> [[Instruction]] -> [[Instruction]]
+    shakeOnce x =
+        filter
+            ( \z -> do
+                let label = head z
+                hasUsage label (concat (filter (\w -> head w /= label) x))
+            )
+    shake x y = do
+        let x' = shakeOnce x y
+        if x' == x then x else shake x' y
+    hasUsage :: Instruction -> [Instruction] -> Bool
+    hasUsage (Label "main") = const True
+    hasUsage (Label ln) = do
+        let l = labelBaseName ln
+        any
+            ( \case
+                Jmp l' -> match l'
+                Jt l' -> match l'
+                Jf l' -> match l'
+                Jz l' -> match l'
+                Jnz l' -> match l'
+                Call l' -> match l'
+                CallLocal l' -> match l'
+                PushPf l' _ -> match l'
+                _ -> False
+            )
+      where
+        match :: String -> Bool
+        match x = (x == ln) || (labelBaseName x == labelBaseName ln)
+    hasUsage _ = const True
+
+    splitLast :: (Eq a) => a -> [a] -> Either [a] ([a], [a])
+    splitLast c' = foldr go (Left [])
+      where
+        go c (Right (f, b)) = Right (c : f, b)
+        go c (Left s)
+            | c' == c = Right ([], s)
+            | otherwise = Left (c : s)
+
+    groupByLabel :: [Instruction] -> [[Instruction]]
+    groupByLabel (x : xs) = case x of
+        Label _ ->
+            let (group, rest) =
+                    span
+                        ( \case
+                            Label _ -> False
+                            _ -> True
+                        )
+                        xs
+             in (x : group) : groupByLabel rest
+        _ -> error "groupByLabel: not implemented"
+    groupByLabel [] = []
+
+    labelBaseName :: String -> String
+    labelBaseName x = do
+        let x' = takeWhile (/= '#') x
+        let x'' = case splitLast ':' x' of
+                Right (_, b) -> b
+                Left _ -> x'
+        x''
+
+optimize :: [Instruction] -> [Instruction]
+optimize = treeShake

lib/VM.hs

@@ -167,7 +167,7 @@ data Data
     | DMap (Data.Map String Data)
     | DTypeQuery String
     | DCPtr WordPtr
-    deriving (Generic)
+    deriving (Generic, Eq)
 
 instance Binary Instruction
 
@@ -193,15 +193,6 @@ instance Show Data where
     show (DCPtr x) = "CPtr " ++ show x
     show (DDouble x) = show x
 
-instance Eq Data where
-    (DInt x) == (DInt y) = x == y
-    (DFloat x) == (DFloat y) = x == y
-    (DString x) == (DString y) = x == y
-    (DBool x) == (DBool y) = x == y
-    (DList x) == (DList y) = x == y
-    (DCPtr x) == (DCPtr y) = x == y
-    x == y = error $ "Cannot eq " ++ show x ++ " and " ++ show y
-
 instance Ord Data where
     (DInt x) `compare` (DInt y) = x `compare` y
     (DFloat x) `compare` (DFloat y) = x `compare` y

tests/IntegrationSpec.hs

@@ -6,8 +6,10 @@ import BytecodeCompiler
     , locateLabel
     )
 import Control.Monad.State (evalStateT)
+import Data.Functor
 import Data.Text qualified
 import Data.Vector qualified as V
+import Optimizer
 import Parser
     ( CompilerFlags (CompilerFlags, verboseMode)
     , parseProgram
@@ -32,7 +34,7 @@ compileAndRun prog = do
     case p of
         Left err -> error $ errorBundlePretty err
         Right program -> do
-            xxx <- evalStateT (compileProgram program) (initCompilerState program)
+            xxx <- evalStateT (compileProgram program) (initCompilerState program) <&> optimize
             let xxxPoint = locateLabel xxx "main"
             -- putStrLn $ printAssembly (V.fromList xxx) True
             vm <- runVMVM $ (initVM (V.fromList xxx)){pc = xxxPoint, breakpoints = [], callStack = [StackFrame{returnAddress = xxxPoint, locals = []}], ioMode = VMBuffer}

tests/OptimizerSpec.hs

@@ -0,0 +1,23 @@
+module OptimizerSpec (spec) where
+
+import Optimizer
+import Test.Hspec
+import VM
+
+spec :: Spec
+spec = do
+    describe "Tree shaker" $ do
+        it "Can tree shake a basic program" $
+            do
+                treeShake
+                    [ Label "hello"
+                    , Push $ DString "Hello, World!"
+                    , Builtin Print
+                    , Ret
+                    , Label "unused"
+                    , Push $ DInt 2
+                    , Ret
+                    , Label "main"
+                    , Call "hello"
+                    ]
+                `shouldBe` [Label "hello", Push $ DString "Hello, World!", Builtin Print, Ret, Label "main", Call "hello"]