Fixes and improvements

6 files changed, 94 insertions, 82 deletions

diff --git a/.gitignore b/.gitignore
index 25c1742..cc3645e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,5 @@ main
 *.hi
 *.asm
 *.exe
+*.ast
+*.succinct
diff --git a/ast.hs b/ast.hs
index 9a4373d..cd7fed2 100644
--- a/ast.hs
+++ b/ast.hs
@@ -14,7 +14,7 @@ newtype Program = Program [Instruction]
 data Instruction
     = IAdd Byte Offset
     | ISet Byte Offset
-    | IMove Offset [(Offset, Byte)]  -- IMove from [(to, multiplier)]
+    | ICopy Offset Offset Byte  -- ICopy from to multiplier
     | ISlide Offset
     | ILoop [Instruction] Offset
     | IInput Offset
@@ -30,18 +30,22 @@ isISet :: Instruction -> Bool
 isISet (ISet _ _) = True
 isISet _ = False
 
-isIMove :: Instruction -> Bool
-isIMove (IMove _ _) = True
-isIMove _ = False
+isICopy :: Instruction -> Bool
+isICopy (ICopy _ _ _) = True
+isICopy _ = False
 
 isISlide :: Instruction -> Bool
 isISlide (ISlide _) = True
 isISlide _ = False
 
+isIStart :: Instruction -> Bool
+isIStart IStart = True
+isIStart _ = False
+
 offsetOf :: Instruction -> Offset
 offsetOf (IAdd _ o) = o
 offsetOf (ISet _ o) = o
-offsetOf (IMove o _) = o
+offsetOf (ICopy o _ _) = o
 offsetOf (ISlide _) = undefined
 offsetOf (ILoop _ _) = undefined
 offsetOf (IInput o) = o
@@ -49,17 +53,16 @@ offsetOf (IOutput o) = o
 offsetOf IStart = 0
 
 astSuccinct :: Program -> String
-astSuccinct (Program inss) = concatMap insSuccinct inss
+astSuccinct (Program inss) = intercalate " " $ map insSuccinct inss
   where
     insSuccinct :: Instruction -> String
     insSuccinct (IAdd v o) =
         let sv = signedByte v
         in (if sv >= 0 then "+" else "") ++ show (signedByte v) ++ ',' : show o
     insSuccinct (ISet v o) = '=' : show (signedByte v) ++ ',' : show o
-    insSuccinct (IMove from tos) =
-        'M' : show from ++ '(' : intercalate "," (map (\(o,m) -> show o ++ '*' : show m) tos) ++ ")"
+    insSuccinct (ICopy from to v) = 'C' : show from ++ ',' : show to ++ ',' : show v
     insSuccinct (ISlide o) = '>' : show o
-    insSuccinct (ILoop inss' off) = "[(" ++ show off ++ ')' : concatMap insSuccinct inss' ++ "]"
+    insSuccinct (ILoop inss' off) = "[(" ++ show off ++ ')' : intercalate " " (map insSuccinct inss') ++ "]"
     insSuccinct (IInput o) = ',' : show o
     insSuccinct (IOutput o) = '.' : show o
     insSuccinct IStart = "$"
diff --git a/compiler.hs b/compiler.hs
index 8a98a2d..eab7058 100644
--- a/compiler.hs
+++ b/compiler.hs
@@ -36,51 +36,48 @@ compile (Program inss) = runLineWriter $ do
     mapM_ compileIns inss
     emit0 epilogue
 
-cursor :: String
-cursor = "rbx"
+cursorReg :: String
+cursorReg = "rbx"
+
+cursorOffset :: Offset -> String
+cursorOffset 0 = cursorReg
+cursorOffset off = cursorReg ++ " + " ++ show off
 
 compileIns :: Instruction -> LineWriter ()
-compileIns (IAdd v off) = emit $ "add byte [" ++ cursor ++ " + " ++ show off ++ "], " ++ show v
-compileIns (ISet v off) = emit $ "mov byte [" ++ cursor ++ " + " ++ show off ++ "], " ++ show v
-compileIns (IMove from []) = do
-    emit $ "mov byte [" ++ cursor ++ " + " ++ show from ++ "], 0"
-compileIns (IMove from [(o,m)]) = do
-    emit $ "mov al, [" ++ cursor ++ " + " ++ show from ++ "]"
-    if m == 1 then return ()
-              else do
-                       emit $ "mov cl, " ++ show m
-                       emit "mul byte cl"
-    emit $ "add [" ++ cursor ++ " + " ++ show o ++ "], al"
-    emit $ "mov byte [" ++ cursor ++ " + " ++ show from ++ "], 0"
-compileIns (IMove from tos) = do
-    emit $ "mov dl, [" ++ cursor ++ " + " ++ show from ++ "]"
-    emit $ "mov byte [" ++ cursor ++ " + " ++ show from ++ "], 0"
-    forM_ tos $ \(o,m) -> do
-        emit "mov al, dl"
-        if m == 1 then return ()
-                  else do
-                           emit $ "mov cl, " ++ show m
-                           emit "mul byte cl"
-        emit $ "add [" ++ cursor ++ " + " ++ show o ++ "], al"
-compileIns (ISlide off) = emit $ "add " ++ cursor ++ ", " ++ show off
+compileIns (IAdd v off) = emit $ "add byte [" ++ cursorOffset off ++ "], " ++ show v
+compileIns (ISet v off) = emit $ "mov byte [" ++ cursorOffset off ++ "], " ++ show v
+compileIns (ICopy _ _ 0) = return ()
+compileIns (ICopy from to m) = do
+    emit $ "mov al, [" ++ cursorOffset from ++ "]"
+    case m of
+        1 -> return ()
+        _ | Just p <- isTwoPower m -> emit $ "shl al, " ++ show p
+          | otherwise -> do
+                             emit $ "mov cl, " ++ show m
+                             emit "mul byte cl"
+    emit $ "add [" ++ cursorOffset to ++ "], al"
+compileIns (ISlide off) = emit $ "add " ++ cursorReg ++ ", " ++ show off
 compileIns (ILoop inss off) = do
     loopid <- genId
-    emit $ "cmp byte [" ++ cursor ++ " + " ++ show off ++ "], 0"
+    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
     emit $ "jz .Laf_" ++ show loopid
     emit0 $ ".Lbd_" ++ show loopid ++ ":"
     mapM_ compileIns inss
-    emit $ "cmp byte [" ++ cursor ++ " + " ++ show off ++ "], 0"
+    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
     emit $ "jnz .Lbd_" ++ show loopid
     emit0 $ ".Laf_" ++ show loopid ++ ":"
 compileIns (IInput off) = do
     emit "call _getchar"
-    emit $ "mov [" ++ cursor ++ " + " ++ show off ++ "], al"
+    emit $ "mov [" ++ cursorOffset off ++ "], al"
 compileIns (IOutput off) = do
     emit "xor edi, edi"
-    emit $ "mov dil, [" ++ cursor ++ " + " ++ show off ++ "]"
+    emit $ "mov dil, [" ++ cursorOffset off ++ "]"
     emit "call _putchar"
 compileIns IStart = return ()
 
+isTwoPower :: Byte -> Maybe Int
+isTwoPower v = findIndex (==v) (take 8 $ iterate (* 2) 1)
+
 
 prologue :: String
 prologue =
diff --git a/interpreter.hs b/interpreter.hs
index 59b9402..4f3f3aa 100644
--- a/interpreter.hs
+++ b/interpreter.hs
@@ -23,10 +23,9 @@ interInss !count !tape !memp allinss@(ins:rest) inp = case ins of
     (ISet value offset) -> do
         MV.write tape (memp + offset) value
         interInss (count + 1) tape memp rest inp
-    (IMove from tos) -> do
+    (ICopy from to mult) -> do
         value <- MV.read tape (memp + from)
-        MV.write tape (memp + from) 0
-        mapM_ (\(offset, multiplier) -> MV.modify tape ((+) (multiplier * value)) (memp + offset)) tos
+        MV.modify tape ((+) (mult * value)) (memp + to)
         interInss (count + 1) tape memp rest inp
     (ISlide offset) -> do
         interInss (count + 1) tape (memp + offset) rest inp
diff --git a/main.hs b/main.hs
index 7d82362..9c33f83 100644
--- a/main.hs
+++ b/main.hs
@@ -1,6 +1,5 @@
 module Main where
 
-import Control.DeepSeq
 import Control.Monad
 import Data.Char
 import System.Environment
@@ -24,18 +23,16 @@ main :: IO ()
 main = do
     args <- getArgs
     when (length args == 0 || length args > 2)
-        $ die "Usage: bfcomphs <source.bf> [source.bf.exe]"
+        $ die "Usage: bfcomphs <source.bf>"
     let fname = head args
-        destfname = if length args == 2 then args !! 1 else fname ++ ".exe"
 
     prog <- readFile fname >>= either die return . parseProgram
 
     -- putStrLn $ astSuccinct prog
     -- print prog
     let opt = optimise prog
-        showopt = force $ show opt
-    when True $ putStrLn $ astSuccinct opt
-    when False $ putStrLn showopt
+    writeFile (fname ++ ".succinct") $ astSuccinct opt
+    writeFile (fname ++ ".ast") $ show opt
 
     case executionMode of
         EMInterpret -> do
@@ -44,4 +41,4 @@ main = do
         EMCompile -> do
             writeFile (fname ++ ".asm") $ compile opt
             callProcess "yasm" ["-f", "macho64", fname ++ ".asm", "-o", fname ++ ".o"]
-            callProcess "gcc" [fname ++ ".o", "-o", destfname]
+            callProcess "gcc" [fname ++ ".o", "-o", fname ++ ".exe"]
diff --git a/optimiser.hs b/optimiser.hs
index 61096b4..2a27a5e 100644
--- a/optimiser.hs
+++ b/optimiser.hs
@@ -1,8 +1,10 @@
 module Optimiser(optimise) where
 
 import Data.List
+import Data.Function
+import qualified Data.Map.Strict as Map
 
--- import Debug.Trace
+import Debug.Trace
 
 import AST
 
@@ -12,11 +14,12 @@ type Optimisation = [Instruction] -> [Instruction]
 
 optimisations :: [Optimisation]
 optimisations =
-    [collectSimilar, nullOps, collectAdds, propagateKnowns, firstSets, propagateSlides, uselessEnd, specialLoops, normaliseOrder, deduplicateSets]
+    [collectSimilar, nullOps, collectAdds, propagateKnowns, firstSets, propagateSlides, uselessEnd, specialLoops, normaliseOrder]
 
 composedOpts :: Optimisation
-composedOpts = foldl1 (.) (reverse optimisations)
+-- composedOpts = foldl1 (.) (reverse optimisations)
 -- composedOpts = foldl1 (.) (map (traceShowId .) $ reverse optimisations)
+composedOpts = foldl1 (.) (map ((\r -> traceShow (take 7 r) r) .) $ reverse optimisations)
 
 optimise :: Program -> Program
 optimise (Program inss) =
@@ -35,12 +38,15 @@ collectSimilar (ins : rest) = ins : collectSimilar rest
 nullOps :: Optimisation
 nullOps [] = []
 nullOps (IAdd 0 _ : rest) = nullOps rest
+nullOps (ICopy _ _ 0 : rest) = nullOps rest
 nullOps (ISlide 0 : rest) = nullOps rest
 nullOps (ILoop inss off : rest) = ILoop (nullOps inss) off : nullOps rest
 nullOps (ins : rest) = ins : nullOps rest
 
 collectAdds :: Optimisation
 collectAdds [] = []
+collectAdds (ICopy from1 to1 m1 : ICopy from2 to2 m2 : rest)
+    | from1 == from2, m1 == m2 = ICopy from1 to1 m1 : collectAdds (ICopy to1 to2 1 : rest)
 collectAdds inss =
     let adds = map (\(IAdd v o) -> (v, o)) $ takeWhile isIAdd inss
         dests = nub $ map snd adds
@@ -53,13 +59,38 @@ collectAdds inss =
 
 propagateKnowns :: Optimisation
 propagateKnowns [] = []
-propagateKnowns (ISet v off : rest) =
-    let pre = takeWhile (\ins -> isIAdd ins || isISet ins || isIMove ins) rest
-        post = drop (length pre) rest
-        relevant = filter ((== off) . offsetOf) pre
-        irrelevant = filter ((/= off) . offsetOf) pre
-        (res, resins) = accumSetAddMove v relevant
-    in ISet res off : resins ++ propagateKnowns (irrelevant ++ post)
+propagateKnowns allinss@(ins : _) | isISet ins || isIStart ins =
+    let defaultZero = isIStart ins
+
+        go :: [Instruction] -> Map.Map Offset Byte -> ([Instruction], Map.Map Offset Byte)
+        go (IStart : rest) mp = go rest mp
+        go (ISet val off : rest) mp = go rest $ Map.insert off val mp
+        go inss@(IAdd val off : rest) mp = case Map.lookup off mp of
+            Nothing -> (inss, mp)
+            Just origval -> go rest $ Map.insert off (origval + val) mp
+        go inss@(ICopy from to mult : rest) mp =
+            if defaultZero
+                then let fromval = maybe 0 id $ Map.lookup from mp
+                         toval = maybe 0 id $ Map.lookup to mp
+                     in go rest $ Map.insert to (toval + mult * fromval) mp
+                else case (Map.lookup from mp, Map.lookup to mp) of
+                         (Nothing, _) -> (inss, mp)
+                         (_, Nothing) -> (inss, mp)
+                         (Just fromval, Just toval) ->
+                             go rest $ Map.insert to (toval + mult * fromval) mp
+        go inss mp = (inss, mp)
+
+        compareSetAdd :: Instruction -> Instruction -> Ordering
+        compareSetAdd (ISet _ _) (IAdd _ _) = LT
+        compareSetAdd (IAdd _ _) (ISet _ _) = GT
+        compareSetAdd i1 i2 = compare (offsetOf i1) (offsetOf i2)
+
+        (after, valuemap) = go allinss Map.empty
+        sets = sortBy compareSetAdd
+                   $ Map.foldlWithKey
+                       (\l off val -> if not defaultZero || val /= 0 then ISet val off : l else l)
+                       [] valuemap
+    in if isIStart ins then IStart : sets ++ propagateKnowns after else sets ++ propagateKnowns after
 propagateKnowns (ILoop inss off : rest) = ILoop (propagateKnowns inss) off : propagateKnowns rest
 propagateKnowns (ins : rest) = ins : propagateKnowns rest
 
@@ -99,25 +130,17 @@ specialLoops (ILoop inss off : rest)
       sum (map (\(IAdd v _) -> v) $ filter ((== off) . offsetOf) inss) == -1 =
         let others = map (\(IAdd v o) -> (v, o)) $ filter ((/= off) . offsetOf) inss
             dests = nub $ map snd others
-            tos = [(d, sum $ map fst $ filter ((== d) . snd) others) | d <- dests]
-        in IMove off tos : specialLoops rest
+            copies = [ICopy off d (sum $ map fst $ filter ((== d) . snd) others) | d <- dests]
+        in copies ++ ISet 0 off : specialLoops rest
 specialLoops (ILoop inss off : rest) = ILoop (specialLoops inss) off : specialLoops rest
 specialLoops (ins : rest) = ins : specialLoops rest
 
 normaliseOrder :: Optimisation
 normaliseOrder [] = []
-normaliseOrder inss =
-    let pre = takeWhile (\ins -> isIAdd ins || isISet ins) inss
-        post = drop (length pre) inss
-    in if null pre
-           then head inss : normaliseOrder (tail inss)
-           else filter isISet pre ++ filter isIAdd pre ++ normaliseOrder post
-
-deduplicateSets :: Optimisation
-deduplicateSets [] = []
-deduplicateSets (IStart : ISet 0 _ : rest) = IStart : deduplicateSets rest
-deduplicateSets (ISet _ o1 : ISet v o2 : rest) | o1 == o2 = deduplicateSets $ ISet v o1 : rest
-deduplicateSets (ins : rest) = ins : deduplicateSets rest
+normaliseOrder (ICopy from1 to1 m1 : ICopy from2 to2 m2 : rest)
+    | m2 < m1, to1 /= from2, to2 /= from1 = ICopy from2 to2 m2 : normaliseOrder (ICopy from1 to1 m1 : rest)
+normaliseOrder (ILoop inss off : rest) = ILoop (normaliseOrder inss) off : normaliseOrder rest
+normaliseOrder (ins : rest) = ins : normaliseOrder rest
 
 
 repeated :: Optimisation -> Optimisation
@@ -128,7 +151,7 @@ incOffsets :: Offset -> [Instruction] -> [Instruction]
 incOffsets _ [] = []
 incOffsets inc (IAdd v off : rest) = IAdd v (off + inc) : incOffsets inc rest
 incOffsets inc (ISet v off : rest) = ISet v (off + inc) : incOffsets inc rest
-incOffsets inc (IMove from tos : rest) = IMove (from + inc) [(o+inc,m) | (o,m) <- tos] : incOffsets inc rest
+incOffsets inc (ICopy from to mult : rest) = ICopy (from + inc) (to + inc) mult : incOffsets inc rest
 incOffsets inc (ISlide off : rest) = ISlide off : incOffsets inc rest
 incOffsets inc (ILoop inss off : rest) = ILoop (incOffsets inc inss) (off + inc) : incOffsets inc rest
 incOffsets inc (IInput off : rest) = IInput (off + inc) : incOffsets inc rest
@@ -141,14 +164,5 @@ accumSetAdd _ (ISet v' _ : rest) = accumSetAdd v' rest
 accumSetAdd acc (IAdd v' _ : rest) = accumSetAdd (acc + v') rest
 accumSetAdd _ _ = unreachable
 
-accumSetAddMove :: Byte -> [Instruction] -> (Byte, [Instruction])
-accumSetAddMove acc [] = (acc, [])
-accumSetAddMove _ (ISet v' _ : rest) = accumSetAddMove v' rest
-accumSetAddMove acc (IAdd v' _ : rest) = accumSetAddMove (acc + v') rest
-accumSetAddMove acc (IMove _ tos : rest) =
-    let (res, resins) = accumSetAddMove 0 rest
-    in (res, [IAdd (m * acc) o | (o,m) <- tos] ++ resins)
-accumSetAddMove _ _ = unreachable
-
 unreachable :: a
 unreachable = error "Unreachable"