1 files changed, 55 insertions, 28 deletions

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