From 9b27451909231ee46ec32ef5adccbe2ab9769903 Mon Sep 17 00:00:00 2001
From: rockofox <ofelix@pm.me>
Date: Tue, 12 Nov 2024 18:33:24 +0100
Subject: [PATCH] Fix trait functions not working as partials (fix #9)

---
 app/Main.hs             | 11 ++++++++---
 lib/BytecodeCompiler.hs | 12 ++++++++++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 52c12f4..959d585 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -10,6 +10,7 @@ import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Identity (Identity (..))
 import Control.Monad.State (StateT (runStateT), evalStateT)
 import Control.Monad.State.Lazy
+import Data.Bool (bool)
 import Data.ByteString.Lazy qualified as B
 import Data.List (intercalate, isPrefixOf)
 import Data.Maybe (fromJust, fromMaybe, isNothing)
@@ -43,6 +44,7 @@ data CmdOptions = Options
     , showTime :: Bool
     , showVersion :: Bool
     , noVerify :: Bool
+    , noOptimize :: Bool
     }
 
 optionsParser :: Options.Applicative.Parser CmdOptions
@@ -69,6 +71,7 @@ optionsParser =
         <*> switch (long "profile" <> short 'p' <> help "Show time spent parsing, compiling and running")
         <*> switch (long "version" <> short 'V' <> help "Show version")
         <*> switch (long "no-verify" <> short 'n' <> help "Don't verify the program")
+        <*> switch (long "no-optimize" <> short 'O' <> help "Don't optimize the program")
 
 prettyPrintExpr :: Expr -> Int -> String
 prettyPrintExpr (DoBlock exprs) i = indent i ++ "DoBlock[\n" ++ intercalate "\n" (map (\x -> prettyPrintExpr x (i + 1)) exprs) ++ "\n" ++ indent i ++ "]"
@@ -121,7 +124,7 @@ parseNoVerify name input compilerFlags = return $ do
 
 main :: IO ()
 main = do
-    Options input output debug verbose emitBytecode runBytecode breakpoints showTime showVersion noVerify <-
+    Options input output debug verbose emitBytecode runBytecode breakpoints showTime showVersion noVerify noOptimize <-
         execParser $
             info
                 (optionsParser <**> helper)
@@ -154,8 +157,10 @@ main = do
                     Right expr -> return expr
 
                 when debug $ putStrLn $ prettyPrintProgram expr
-                potentiallyTimedOperation "Compilation" showTime $ do
-                    optimize <$> evalStateT (BytecodeCompiler.compileProgram expr) (BytecodeCompiler.initCompilerState expr)
+                potentiallyTimedOperation "Compilation" showTime $
+                    do
+                        bool id optimize (not noOptimize)
+                        <$> evalStateT (BytecodeCompiler.compileProgram expr) (BytecodeCompiler.initCompilerState expr)
             else do
                 bytecode <- inputFileBinary input
                 return $ VM.fromBytecode bytecode
diff --git a/lib/BytecodeCompiler.hs b/lib/BytecodeCompiler.hs
index b8481ba..4c9a4d5 100644
--- a/lib/BytecodeCompiler.hs
+++ b/lib/BytecodeCompiler.hs
@@ -2,15 +2,18 @@ module BytecodeCompiler where
 
 import AST (zeroPosition)
 import AST qualified as Parser.Type (Type (Unknown))
+import Control.Arrow ((>>>))
 import Control.Monad (when, (>=>))
 import Control.Monad.Loops (firstM)
 import Control.Monad.State (MonadIO (liftIO), StateT, gets, modify)
 import Data.Bifunctor (second)
+import Data.Function ((&))
 import Data.Functor ((<&>))
 import Data.List (elemIndex, find, inits, intercalate)
 import Data.List.Split qualified
 import Data.Map qualified
 import Data.Maybe (fromJust, isJust, isNothing)
+import Data.String
 import Data.Text (isPrefixOf, splitOn)
 import Data.Text qualified
 import Data.Text qualified as T
@@ -312,7 +315,9 @@ compileExpr (Parser.FuncCall funcName args _) = do
                                 return $
                                     args'
                                         ++ [PushPf (funame fun) (length args')]
-                Nothing ->
+                Nothing -> do
+                    -- traceM $ "Looking for funcDec " ++ funcName
+                    -- gets funcDecs >>= \x -> traceM $ "\t" ++ show (map Parser.name x)
                     concatMapM compileExpr args >>= \args' ->
                         return $
                             args'
@@ -381,7 +386,10 @@ compileExpr (Parser.Var x _) = do
     curCon <- gets currentContext
     externals' <- gets externals
 
-    let fun = any ((== x) . baseName) functions' || any ((== curCon ++ "@" ++ x) . baseName) functions'
+    let fun =
+            any ((== x) . baseName) functions'
+                || any ((== curCon ++ "@" ++ x) . baseName) functions'
+                || any ((== x) . ((Data.Text.unpack . last . splitOn "::") . fromString . baseName)) functions'
     if fun || x `elem` internalFunctions || x `elem` map (\f -> f.name) externals'
         then compileExpr (Parser.FuncCall x [] zeroPosition)
         else return [LLoad x]