Initial working optimization stage

8 files changed, 77 insertions, 36 deletions

diff --git a/app/Main.hs b/app/Main.hs
index 8731b7c..47d1bd2 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -24,6 +24,10 @@ args =
     <$> (mode <|> pure ArgEval)
     <*> switch (long "yolo" <> short 'y' <> help "Don't run tests")
     <*> switch (long "verbose" <> short 'v' <> help "Increase verbosity")
+    <*> switch
+          (long "optimize" <> short 'O' <> help
+            "Optimize program (abstraction of duplicated terms)"
+          )
     <*> strOption
           (long "target" <> short 't' <> metavar "TARGET" <> value "" <> help
             "Compile to target using BLoC and BLoCade"
diff --git a/benchmarks/standalone/optimization.bruijn b/benchmarks/standalone/optimization.bruijn
new file mode 100644
index 0000000..564c95f
--- /dev/null
+++ b/benchmarks/standalone/optimization.bruijn
@@ -0,0 +1,3 @@
+:import std/Math .
+
+main [(fac (+100)) + (fac (+100))]
diff --git a/docs/wiki_src/coding/laziness.md b/docs/wiki_src/coding/laziness.md
index 20d1205..e3d642f 100644
--- a/docs/wiki_src/coding/laziness.md
+++ b/docs/wiki_src/coding/laziness.md
@@ -38,11 +38,11 @@ Laziness can (in some cases) produce huge performance boosts. For
 example:
 
 ``` bruijn
-# 11 seconds
-:time (+10) ** (+500)
+# 10 seconds
+:time (+10) ** (+1000)
 
-# 0.1 seconds
-:time ((+10) ** (+500)) =? (+400)
+# 0.02 seconds
+:time ((+10) ** (+1000)) =? (+42)
 ```
 
 This works because a ternary number is just a concatenation of trits
diff --git a/docs/wiki_src/coding/recursion.md b/docs/wiki_src/coding/recursion.md
index 89e71c8..173cd40 100644
--- a/docs/wiki_src/coding/recursion.md
+++ b/docs/wiki_src/coding/recursion.md
@@ -23,7 +23,7 @@ Say we want a function `g`{.bruijn} to be able to call itself. With the
 With this equivalence, `g`{.bruijn} is able to call itself since its
 outer argument is the initial function again.
 
-Example for using `y`{.bruijn} to find the factorial of 2:
+Example for using `y`{.bruijn} to find the factorial of 3:
 
 ``` bruijn
 # here, `1` is the induced outer argument `(y g)`
@@ -57,7 +57,8 @@ suggestions.
 For solving mutual recurrence relations, you can use the *variadic
 fixed-point combinator* `y*`{.bruijn} from
 [`std/List`](/std/List.bruijn.html). This combinator produces all the
-fixed points of a function as an iterable [list](data-structures.md).
+fixed points of given functions as an iterable
+[list](data-structures.md).
 
 Example `even?`{.bruijn}/`odd?`{.bruijn} implementation using
 `y*`{.bruijn}:
diff --git a/docs/wiki_src/technical/performance.md b/docs/wiki_src/technical/performance.md
index 709affe..ce3cca0 100644
--- a/docs/wiki_src/technical/performance.md
+++ b/docs/wiki_src/technical/performance.md
@@ -11,6 +11,7 @@ comparable as possible:
 -   Bruijn uses efficient data structures by default. For example, for
     nary numbers we use results of Torben Mogensens investigations (as
     described in [number/byte encodings](../coding/data-structures.md)).
+-   Bruijn has a `-O` flag that enables abstraction of duplicated terms
 -   The lambda calculus optimizers
     [BLoC](https://github.com/marvinborner/bloc) and
     [BLoCade](https://github.com/marvinborner/blocade) are directly
diff --git a/src/Eval.hs b/src/Eval.hs
index d331e86..0a5e8c7 100644
--- a/src/Eval.hs
+++ b/src/Eval.hs
@@ -419,7 +419,7 @@ evalFileConf conf = do
     Nothing -> print $ ContextualError (UndefinedIdentifier entryFunction)
                                        (Context "" (_nicePath conf))
     Just EnvDef { _exp = e } -> do
-      red <- optimizedReduce conf (Application e arg)
+      red <- optimizedReduce conf { _hasArg = True } (Application e arg)
       showResult red (Environment env)
 
 exec :: EvalConf -> (String -> IO (Either IOError a)) -> (a -> String) -> IO ()
@@ -429,7 +429,7 @@ exec conf rd conv = do
   case f of
     Left  exception -> print (exception :: IOError)
     Right f'        -> do
-      red <- optimizedReduce conf (Application e arg)
+      red <- optimizedReduce conf { _hasArg = True } (Application e arg)
       showResult red (Environment M.empty)
       where e = fromBinary $ conv f'
 
diff --git a/src/Helper.hs b/src/Helper.hs
index b4b48e1..edd9b09 100644
--- a/src/Helper.hs
+++ b/src/Helper.hs
@@ -187,6 +187,7 @@ data Args = Args
   { _argMode     :: ArgMode
   , _argNoTests  :: Bool
   , _argVerbose  :: Bool
+  , _argOptimize :: Bool
   , _argToTarget :: String
   , _argReducer  :: String
   , _argPath     :: Maybe String
@@ -196,11 +197,13 @@ data EvalConf = EvalConf
   { _isRepl    :: Bool
   , _isVerbose :: Bool
   , _evalTests :: Bool
+  , _optimize  :: Bool
   , _nicePath  :: String
   , _path      :: String
   , _evalPaths :: [String]
   , _toTarget  :: String
   , _reducer   :: String
+  , _hasArg    :: Bool
   }
 
 newtype ExpFlags = ExpFlags
@@ -228,11 +231,13 @@ argsToConf :: Args -> EvalConf
 argsToConf args = EvalConf { _isRepl    = isNothing $ _argPath args
                            , _isVerbose = _argVerbose args
                            , _evalTests = not $ _argNoTests args
+                           , _optimize  = _argOptimize args
                            , _path      = path
                            , _nicePath  = path
                            , _evalPaths = []
                            , _toTarget  = _argToTarget args
                            , _reducer   = _argReducer args
+                           , _hasArg    = False
                            }
   where path = fromMaybe "" (_argPath args)
 
diff --git a/src/Optimizer.hs b/src/Optimizer.hs
index e56f12d..e7f56ab 100644
--- a/src/Optimizer.hs
+++ b/src/Optimizer.hs
@@ -1,8 +1,10 @@
 -- MIT License, Copyright (c) 2024 Marvin Borner
+-- TODO: This currently only really helps in edge-cases (see benchmarks) -> it's disabled by default
 module Optimizer
   ( optimizedReduce
   ) where
 
+import           Data.List                      ( tails )
 import qualified Data.Map                      as M
 import           Helper
 import           Reducer
@@ -29,7 +31,7 @@ constructTree = go [] M.empty
   go p m e@(Abstraction t) = do
     let m' = go (D : p) m t
     M.insertWith (++) e [p] m'
-  go p m e@(Bruijn i) = M.insertWith (++) e [p] m
+  go p m e@(Bruijn _) = M.insertWith (++) e [p] m
   go _ _ _            = invalidProgramState
 
 isClosed :: Expression -> Bool
@@ -42,7 +44,7 @@ isClosed = go 0
 
 -- (kinda arbitrary)
 isGood :: Expression -> Bool
-isGood = go 10
+isGood = go (10 :: Int)
  where
   go 0 _                               = True
   go i (Application (Abstraction l) r) = go (i - 2) l || go (i - 2) r
@@ -50,19 +52,22 @@ isGood = go 10
   go i (Abstraction t                ) = go (i - 1) t
   go _ _                               = False
 
--- if expression has a parent that appears more often
+firstLast :: [a] -> [a]
+firstLast xs@(_ : _) = tail (init xs)
+firstLast _          = []
+
+-- if expression has a parent that is also "good"
 preferParent :: Expression -> [Path] -> Tree -> Bool
-preferParent e ps t = do
-  let len = length ps
-  any
-    (\p -> any
-      (\p' -> case resolvePath e p' of
-        Just r  -> length (M.findWithDefault [] r t) >= len
-        Nothing -> False
-      )
-      [R : p, D : p, L : p]
+preferParent e ps t = any
+  (\p -> any
+    (\p' -> not (null p') && case resolvePath e (reverse p') of
+        -- Just r  -> length (M.findWithDefault [] r t) >= length ps
+      Just r  -> M.member r t
+      Nothing -> False
     )
-    ps
+    (firstLast $ tails p)
+  )
+  ps
 
 commonPath :: Path -> Path -> Path
 commonPath p1 p2 = go (reverse p1) (reverse p2)
@@ -72,13 +77,36 @@ commonPath p1 p2 = go (reverse p1) (reverse p2)
   go (x : xs) (y : ys) | x == y    = x : go xs ys
                        | otherwise = []
 
--- inject :: Expression -> Path -> Expression -> [Path] -> Expression
--- inject i [] e ps = Application (Abstraction (subst ? (incv e))) i
--- inject i [L : p] (Application l r) ps = Application (inject i p l ps) r
--- inject i [R : p] (Application l r) ps = Application l (inject i p r ps)
--- inject i [D : p] (Abstraction t  ) ps = Abstraction (inject i p t ps)
--- inject _ _       _                 _  = invalidProgramState
+incv :: Expression -> Expression
+incv = go 0
+ where
+  go i (Bruijn j       ) = Bruijn $ if i <= j then j + 1 else j
+  go i (Application a b) = Application (go i a) (go i b)
+  go i (Abstraction t  ) = Abstraction $ go (i + 1) t
+  go _ _                 = invalidProgramState
+
+-- term to index
+subst :: Expression -> Expression -> Expression
+subst = go 0
+ where
+  go i n h = if n == h
+    then Bruijn i
+    else case h of
+      (Application a b) -> Application (go i n a) (go i n b)
+      (Abstraction t  ) -> Abstraction $ go (i + 1) n t
+      t                 -> t
+
+inject :: Expression -> Path -> Expression -> Expression
+inject i [] e = Application (Abstraction (subst i (incv e))) i
+inject i (L : p) (Application l r) = Application (inject i p l) r
+inject i (R : p) (Application l r) = Application l (inject i p r)
+inject i (D : p) (Abstraction t  ) = Abstraction (inject i p t)
+
+-- TODO: Could this produce wrong (or redundant) results?
+-- Optimally, `common` below would be topologically sorted so this can't happen
+inject _ _       e                 = e
 
+-- TODO: Also abstract open terms (using max index within term)
 optimize :: Expression -> IO Expression
 optimize e = do
   let tree = constructTree e
@@ -86,14 +114,13 @@ optimize e = do
         M.filterWithKey (\k ps -> isClosed k && isGood k && length ps > 1) tree
   -- TODO: simulated annealing on every closed term even if not good or ==1
   let filtered' =
-        M.filterWithKey (\k ps -> not $ preferParent e ps filtered) filtered
-  print $ (\(k, p) -> foldl1 commonPath p) <$> M.toList filtered'
-  -- inject t (take (length commonPath) ps) e ps -- oder so
-  pure e
--- optimize e = constructTree e
+        M.filterWithKey (\_ ps -> not $ preferParent e ps filtered) filtered
+  let common = (\(k, p) -> (k, foldl1 commonPath p)) <$> M.toList filtered'
+  pure $ foldl (\e' (t, p) -> inject t p e') e common
 
--- TODO: enable optimizer with flag
 optimizedReduce :: EvalConf -> Expression -> IO Expression
-optimizedReduce conf e = do
-  -- optimize e
-  reduce conf e
+optimizedReduce conf@EvalConf { _optimize = False } e = reduce conf e
+optimizedReduce conf@EvalConf { _hasArg = True } (Application e arg) = do
+  e' <- optimize e
+  reduce conf (Application e' arg)
+optimizedReduce conf e = optimize e >>= reduce conf