improve fixpoint algorithm

lib/src/Control/Arrow/Transformer/FixpointCache.hs

@@ -8,13 +8,12 @@
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE MonoLocalBinds #-}
 {-# LANGUAGE CPP #-}
-{-# OPTIONS_GHC -DTRACE #-}
+-- {-# OPTIONS_GHC -DTRACE #-}
 module Control.Arrow.Transformer.FixpointCache(CacheArrow,runCacheArrow,runCacheArrow',liftCache) where
 
 import           Prelude hiding (id,(.),lookup)
-import           Data.Function (fix)
 
-import           Control.Arrow
+import           Control.Arrow hiding (loop)
 import           Control.Arrow.Class.Fix
 import           Control.Arrow.Utils
 import           Control.Category
@@ -30,112 +29,87 @@ import           Debug.Trace
 import           Text.Printf
 #endif
 
-newtype CacheArrow a b x y = CacheArrow (((Store a b,Store a b),x) -> (Store a b,y))
+newtype CacheArrow a b x y = CacheArrow ((Store a b,x) -> (Store a b,y))
 
 runCacheArrow :: CacheArrow a b x y -> (x -> y)
 runCacheArrow f = runCacheArrow' f >>^ snd
 
 runCacheArrow' :: CacheArrow a b x y -> (x -> (Store a b,y))
-runCacheArrow' (CacheArrow f) = (\x -> ((S.empty,S.empty),x)) ^>> f
+runCacheArrow' (CacheArrow f) = (\x -> ((S.empty),x)) ^>> f
 
 liftCache :: (x -> y) -> CacheArrow a b x y
-liftCache f = CacheArrow ((\((_,o),x) -> (o,x)) ^>> second f)
+liftCache f = CacheArrow ((\(o,x) -> (o,x)) ^>> second f)
 
 instance Category (CacheArrow i o) where
   id = liftCache id
-  CacheArrow f . CacheArrow g = CacheArrow $ proc ((i,o),x) -> do
-    (o',y) <- g -< ((i,o),x)
-    f -< ((i,o'),y)
+  CacheArrow f . CacheArrow g = CacheArrow $ proc (o,x) -> do
+    (o',y) <- g -< (o,x)
+    f -< (o',y)
 
 instance Arrow (CacheArrow i o) where
   arr f = liftCache (arr f)
-  first (CacheArrow f) = CacheArrow $ (\((i,o),(x,y)) -> (((i,o),x),y)) ^>> first f >>^ (\((o,x'),y) -> (o,(x',y)))
-  second (CacheArrow f) = CacheArrow $ (\((i,o),(x,y)) -> (x,((i,o),y))) ^>> second f >>^ (\(x,(o,y')) -> (o,(x,y')))
+  first (CacheArrow f) = CacheArrow $ (\(o,(x,y)) -> ((o,x),y)) ^>> first f >>^ (\((o,x'),y) -> (o,(x',y)))
+  second (CacheArrow f) = CacheArrow $ (\(o,(x,y)) -> (x,(o,y))) ^>> second f >>^ (\(x,(o,y')) -> (o,(x,y')))
 
 instance ArrowChoice (CacheArrow i o) where
-  left (CacheArrow f) = CacheArrow $ (\((i,o),e) -> injectRight (o,injectLeft ((i,o),e))) ^>> left f >>^ eject
-  right (CacheArrow f) = CacheArrow $ (\((i,o),e) -> injectRight ((i,o),injectLeft (o,e))) ^>> right f >>^ eject
+  left (CacheArrow f) = CacheArrow $ (\(o,e) -> injectRight (o,injectLeft (o,e))) ^>> left f >>^ eject
+  right (CacheArrow f) = CacheArrow $ (\(o,e) -> injectRight (o,injectLeft (o,e))) ^>> right f >>^ eject
 
 instance ArrowApply (CacheArrow i o) where
-  app = CacheArrow $ (\(io,(CacheArrow f,x)) -> (f,(io,x))) ^>> app
+  app = CacheArrow $ (\(o,(CacheArrow f,x)) -> (f,(o,x))) ^>> app
 
 #ifdef TRACE
 instance (Show x, Show y, Eq x, Hashable x, LowerBounded y, Complete y)
   => ArrowFix x y (CacheArrow x y) where
-  fixA f = trace (printf "fixA fact") $ proc x -> do
-    old <- getOutCache -< ()
-    setOutCache -< bottom
-    y <- localInCache (trace (printf "\tfix (memoize . fact)") $ fix (memoize . f)) -< (old,x)
-    new <- getOutCache -< ()
-    if new ⊑ old -- We are in the reductive set of `f` and have overshot the fixpoint
-    then returnA -< y
-    else fixA f -< x
-
-memoize :: (Show x, Show y, Eq x, Hashable x, LowerBounded y, Complete y) => CacheArrow x y x y -> CacheArrow x y x y
-memoize f = proc x -> do
-  m <- lookupOutCache -< trace (printf "\t\tmemoize -< %s" (show x)) x
-  case m of
-    Just y -> do
-      returnA -< trace (printf "\t\t%s <- memoize -< %s" (show y) (show x)) y
-    Nothing -> do
-      yOld <- lookupInCache -< x
-      writeOutCache -< trace (printf "\t\tout(%s) := %s" (show x) (show (fromMaybe bottom yOld))) (x, fromMaybe bottom yOld)
-      y <- f -< trace (printf "\t\tfact -< %s" (show x)) x
-      yCached <- lookupOutCache -< x
-      updateOutCache -< trace (printf "\t\tout(%s) := %s ⊔ %s = %s" (show x) (show (fromJust yCached)) (show y) (show (fromJust yCached ⊔ y)))
-                         (trace (printf "\t\t%s <- fact -< %s" (show y) (show x))
-                        (x, y))
-      returnA -< trace (printf "\t\t%s <- memoize -< %s" (show y) (show x)) y
-#elif
-
+  fixA f = proc x -> do
+    m <- lookupCache -< x
+    case m of
+      Just y -> do
+        returnA -< trace (printf "%s <- memoize -< %s\t(HIT)" (show y) (show x)) y
+      Nothing -> do
+        updateCache -< trace (printf "memoize -< %s\t\t(MISS)\nout(%s) := ⊥" (show x) (show x)) (x, bottom)
+        y <- loop -< x
+        returnA -< trace (printf "%s <- memoize -< %s" (show y) (show x)) y
+      where
+        loop = proc x -> do
+          yOld <- fromJust ^<< lookupCache -< x
+          y <- f (fixA f) -< trace (printf "%s -< %s" name (show x)) x
+          yNew <- updateCache -< trace (printf "%s <- %s -< %s" (show y) name (show x)) (x, y)
+          if trace (printf "out(%s) := %s ⊔ %s = %s" (show x) (show yOld) (show y) (show yNew)) $
+             yNew ⊑ yOld
+          then returnA -< trace (printf "%s ⊑ %s\t\t\t(RETURN)" (show yNew) (show yOld)) yNew
+          else loop -< trace (printf "%s ⋢ %s\t\t\t(REPEAT)" (show yNew) (show yOld)) x
+        name = "fact"
+#else
 instance (Eq x, Hashable x, LowerBounded y, Complete y)
   => ArrowFix x y (CacheArrow x y) where
   fixA f = proc x -> do
-    old <- getOutCache -< ()
-    setOutCache -< bottom
-    y <- localInCache (fix (memoize . f)) -< (old,x)
-    new <- getOutCache -< ()
-    if (new ⊑ old) -- We are in the reductive set of `f` and have overshot the fixpoint
-    then returnA -< y
-    else fixA f -< x
-
-memoize :: (Eq x, Hashable x, LowerBounded y, Complete y) => CacheArrow x y x y -> CacheArrow x y x y
-memoize f = proc x -> do
-  m <- lookupOutCache -< x
-  case m of
-    Just y -> do
-      returnA -< y
-    Nothing -> do
-      yOld <- lookupInCache -< x
-      writeOutCache -< (x, fromMaybe bottom yOld)
-      y <- f -< x
-      updateOutCache -< (x, y)
-      returnA -< y
+    m <- lookupCache -< x
+    case m of
+      Just y -> returnA -< y
+      Nothing -> do
+        updateCache -< (x, bottom)
+        loop -< x
+      where
+        loop = proc x -> do
+          yOld <- fromJust ^<< lookupCache -< x
+          y <- f (fixA f) -< x
+          yNew <- updateCache -< (x, y)
+          if yNew ⊑ yOld
+          then returnA -< yNew
+          else loop -< x
 #endif
 
-lookupOutCache :: (Eq x, Hashable x) => CacheArrow x y x (Maybe y)
-lookupOutCache = CacheArrow $ \((_,o),x) -> (o,S.lookup x o)
-
-lookupInCache :: (Eq x, Hashable x) => CacheArrow x y x (Maybe y)
-lookupInCache = CacheArrow $ \((i,o),x) -> (o,S.lookup x i)
-
-writeOutCache :: (Eq x, Hashable x) => CacheArrow x y (x,y) ()
-writeOutCache = CacheArrow $ \((_,o),(x,y)) -> (S.insert x y o,())
-
-getOutCache :: CacheArrow x y () (Store x y)
-getOutCache = CacheArrow $ (\((_,o),()) -> (o,o))
-
-setOutCache :: CacheArrow x y (Store x y) ()
-setOutCache = CacheArrow $ (\((_,_),o) -> (o,()))
-
-localInCache :: CacheArrow x y x y -> CacheArrow x y (Store x y,x) y
-localInCache (CacheArrow f) = CacheArrow (\((_,o),(i,x)) -> f ((i,o),x))
+lookupCache :: (Eq x, Hashable x) => CacheArrow x y x (Maybe y)
+lookupCache = CacheArrow $ \(o,x) -> (o,S.lookup x o)
 
-updateOutCache :: (Eq x, Hashable x, Complete y) => CacheArrow x y (x,y) ()
-updateOutCache = CacheArrow $ \((_,o),(x,y)) -> (S.insertWith (⊔) x y o,())
+updateCache :: (Eq x, Hashable x, Complete y) => CacheArrow x y (x,y) y
+updateCache = CacheArrow $ \(o,(x,y)) -> case S.lookup x o of
+  Just y' -> let y'' = y ⊔ y' in (S.insert x y'' o,y'')
+  Nothing -> (S.insert x y o,y)
 
-deriving instance PreOrd (((Store a b,Store a b),x) -> (Store a b,y)) => PreOrd (CacheArrow a b x y)
-deriving instance Complete (((Store a b,Store a b),x) -> (Store a b,y)) => Complete (CacheArrow a b x y)
-deriving instance CoComplete (((Store a b,Store a b),x) -> (Store a b,y)) => CoComplete (CacheArrow a b x y)
-deriving instance LowerBounded (((Store a b,Store a b),x) -> (Store a b,y)) => LowerBounded (CacheArrow a b x y)
-deriving instance UpperBounded (((Store a b,Store a b),x) -> (Store a b,y)) => UpperBounded (CacheArrow a b x y)
+deriving instance PreOrd ((Store a b,x) -> (Store a b,y)) => PreOrd (CacheArrow a b x y)
+deriving instance Complete ((Store a b,x) -> (Store a b,y)) => Complete (CacheArrow a b x y)
+deriving instance CoComplete ((Store a b,x) -> (Store a b,y)) => CoComplete (CacheArrow a b x y)
+deriving instance LowerBounded ((Store a b,x) -> (Store a b,y)) => LowerBounded (CacheArrow a b x y)
+deriving instance UpperBounded ((Store a b,x) -> (Store a b,y)) => UpperBounded (CacheArrow a b x y)

lib/test/FixpointCacheSpec.hs

@@ -28,7 +28,7 @@ main = hspec spec
 
 type Cache x y = CacheArrow x y x y
 type ErrorCache x y = ErrorArrow () (CacheArrow x (Error () y)) x y
-type StateCache x y = StateArrow Int (CacheArrow (Int,x) (Int,y)) x y
+type StateCache x y = StateArrow (InfiniteNumber Int) (CacheArrow (InfiniteNumber Int,x) (InfiniteNumber Int,y)) x y
 type IV = Interval (InfiniteNumber Int)
 
 spec :: Spec
@@ -70,10 +70,10 @@ spec = do
           _ -> f -< (n-1)
     in it "should fail, but update the fixpoint cache" $
          runCacheArrow' (runErrorArrow (fixA recurseFail)) 5
-            `shouldBe` (S.fromList [(n,Error ()) | n <- [0..4]], Error ())
+            `shouldBe` (S.fromList [(n,Error ()) | n <- [0..5]], Error ())
 
   describe "the analysis of a stateful program" $
-    let timesTwo :: StateCache Int () -> StateCache Int ()
+    let timesTwo :: StateCache (InfiniteNumber Int) () -> StateCache (InfiniteNumber Int) ()
         timesTwo f = proc n -> case n of
           0 -> returnA -< ()
           _ -> do
@@ -85,7 +85,7 @@ spec = do
 
     in it "should cache the state of the program" $
          runCacheArrow' (runStateArrow (fixA timesTwo)) (0,5)
-           `shouldBe` (S.fromList [((n,5-n),(10-n,())) | n <- [1..5]],(10,()))
+           `shouldBe` (S.fromList [((Number n,5- Number n),(10-Number n,())) | n <- [0..5]],(10,()))
   where
 
     ifEq :: (Num n, Ord n, ArrowChoice c, Complete (c (Interval n, Interval n) (Interval n)), Complete (c (Interval n,(Interval n, Interval n)) (Interval n))) => n -> c (Interval n) (Interval n) -> c (Interval n) (Interval n) -> c (Interval n) (Interval n)