Add addaptive algorithm

rtree/src/Control/Monad/IRTree.hs

+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TypeOperators #-}
 
 module Control.Monad.IRTree (
   -- * IRTree
@@ -11,16 +14,18 @@ module Control.Monad.IRTree (
   extract,
   probe,
   reduce,
-  reduceExp,
-  reduceFib,
+  reduceAdaptive,
 
   -- * IRTreeT
   IRTreeT,
   probeT,
   extractT,
   reduceT,
-  reduceFibT,
-  reduceExpT,
+  reduceAdaptiveT,
+
+  -- * helpers
+  binsearch,
+  expsearch,
 
   -- * Re-exports
   module Control.Monad.Reader,
@@ -32,6 +37,7 @@ import Control.Monad.Reader
 import Control.Monad.Reduce
 import Data.Bits
 import Data.Foldable
+import Data.Function
 import Data.Functor
 import Data.Functor.Identity
 import Data.RPath
@@ -56,123 +62,104 @@ extractT :: (Functor m) => IRTreeT l m i -> m i
 extractT (IRTreeT m) = fmap (\(i, _, _) -> i) (runRWST m "" 0)
 {-# INLINE extractT #-}
 
-probe :: IRTree l i -> RPath -> (i, [(Bool, l)])
+probe :: IRTree l i -> RPath -> (i, [AnnotatedChoice l], Int)
 probe t pth = runIdentity $ probeT t pth
 {-# INLINE probe #-}
 
-probeT :: (Functor m) => IRTreeT l m i -> RPath -> m (i, [(Bool, l)])
+{- | Probe the IRTree at a path RPath, and return the input, a list of annotated choices,
+and the number of guesses made (can be negative to indicate that not all choices in
+the path was used.
+-}
+probeT :: (Functor m) => IRTreeT l m i -> RPath -> m (i, [AnnotatedChoice l], Int)
 probeT (IRTreeT m) pth =
-  runRWST m pth 0 <&> \(i, _, l) ->
-    (i, zip (toChoiceList pth <> repeat False) (appEndo l []))
+  runRWST m pth 0 <&> \(i, x, l) ->
+    (i, annotateChoices pth (appEndo l []), x - numberOfChoices pth)
 {-# INLINE probeT #-}
 
 reduce
-  :: (Monad m)
-  => ([(Bool, l)] -> i -> m Bool)
+  :: (Monad m, m ~ IO, Show l)
+  => ([AnnotatedChoice l] -> i -> m Bool)
   -> IRTree l i
-  -> m (Maybe i)
+  -> m i
 reduce = reduceT (pure . runIdentity)
 {-# INLINE reduce #-}
 
 -- | Interpreted reduction with an m base monad
 reduceT
-  :: (Monad m, Functor t)
+  :: (Monad m, Functor t, m ~ IO, Show l)
   => (forall a. t a -> m a)
-  -> ([(Bool, l)] -> i -> m Bool)
+  -> ([AnnotatedChoice l] -> i -> m Bool)
   -> IRTreeT l t i
-  -> m (Maybe i)
+  -> m i
 reduceT lift_ p rt = do
-  (i, l) <- lift_ (probeT rt "")
-  t <- p l i
-  if t
-    then Just <$> go Seq.empty
-    else pure Nothing
- where
-  go pth = do
+  Seq.empty & fix \rec sq -> do
     -- Try to run the true branch.
-    let pth' = fromChoiceList (toList (pth Seq.|> True))
-    (i, l) <- lift_ $ probeT rt pth'
-    if length l >= numberOfChoices pth'
-      then do
-        t <- p l i
-        go (pth Seq.|> t)
-      else pure i
+    (i, l, left) <- _probe (sq Seq.|> True)
+    p l i >>= \case
+      -- If predicate is true, and there is choices left
+      True | left > 0 -> rec (sq Seq.|> True)
+      -- If predicate is false (and stable)
+      False | left >= 0 -> rec (sq Seq.|> False)
+      _ow -> pure i
+ where
+  _probe sq = lift_ . probeT rt . fromChoiceList $ toList sq
 {-# INLINE reduceT #-}
 
-reduceExp
-  :: (Monad m)
-  => ([(Bool, l)] -> i -> m Bool)
-  -> IRTree l i
-  -> m (Maybe i)
-reduceExp = reduceExpT (pure . runIdentity)
-{-# INLINE reduceExp #-}
-
-reduceFib
-  :: (Monad m)
-  => ([(Bool, l)] -> i -> m Bool)
+reduceAdaptive
+  :: (Monad m, m ~ IO, Show i, Show l)
+  => ([AnnotatedChoice l] -> i -> m Bool)
   -> IRTree l i
-  -> m (Maybe i)
-reduceFib = reduceFibT (pure . runIdentity)
-{-# INLINE reduceFib #-}
+  -> m i
+reduceAdaptive = reduceAdaptiveT (pure . runIdentity)
+{-# INLINE reduceAdaptive #-}
 
--- | Interpreted reduction with an m base monad, and running in expoential mode.
-reduceExpT
-  :: (Monad m, Functor t)
+-- | Interpreted reduction with an m base monad, but using exponential search.
+reduceAdaptiveT
+  :: (Monad m, Functor t, IO ~ m, Show i, Show l)
   => (forall a. t a -> m a)
   -- ^ a lift of monad m into t (normally @id@ or @lift@)
-  -> ([(Bool, l)] -> i -> m Bool)
+  -> ([AnnotatedChoice l] -> i -> m Bool)
   -> IRTreeT l t i
-  -> m (Maybe i)
-reduceExpT lift_ p rt = do
-  (i, l) <- lift_ (probeT rt "")
-  t <- p l i
-  if t
-    then Just <$> go 0 Seq.empty
-    else pure Nothing
+  -> m i
+reduceAdaptiveT lift_ p rt = do
+  Seq.empty & fix \rec !sq -> do
+    (d', c) <- expsearch \d -> do
+      let sq' = sq <> Seq.replicate d True
+      (i, l, left) <- _probe sq'
+      t <- p l i
+      pure (t, left)
+    if c
+      then rec (sq <> Seq.replicate d' True <> Seq.singleton False)
+      else do
+        (i, _, _) <- _probe (sq <> Seq.replicate d' True)
+        pure i
  where
-  go n sq = do
-    let depth = shiftL 1 n
-    let sq' = sq <> Seq.replicate depth True
-    let pth' = fromChoiceList (toList sq')
-    (i, l) <- lift_ (probeT rt pth')
-    if length l >= numberOfChoices pth' - depth + 1
-      then do
-        t <- p l i
-        if t
-          then go (n + 1) sq'
-          else case n of
-            0 -> go 0 (sq Seq.|> False)
-            n' -> go (n' - 1) sq
-      else pure i
-{-# INLINE reduceExpT #-}
-
--- | Interpreted reduction with an m base monad, and running in fibonacci mode.
-reduceFibT
-  :: (Monad m, Functor t)
-  => (forall a. t a -> m a)
-  -- ^ a lift of monad m into t (normally @id@ or @lift@)
-  -> ([(Bool, l)] -> i -> m Bool)
-  -> IRTreeT l t i
-  -> m (Maybe i)
-reduceFibT lift_ p rt = do
-  (i, l) <- lift_ (probeT rt "")
-  t <- p l i
-  if t
-    then Just <$> go 1 1 Seq.empty
-    else pure Nothing
+  _probe sq = lift_ . probeT rt . fromChoiceList $ toList sq
+
+expsearch :: (Monad m) => (Int -> m (Bool, Int)) -> m (Int, Bool)
+expsearch p =
+  1 & fix \rec !d -> do
+    p d >>= \case
+      (True, left)
+        | left < 0 -> pure (d + left, False)
+        | otherwise -> rec (d `shiftL` 1)
+      (False, left) -> do
+        d' <-
+          binsearch
+            (fmap (fmap fst) p)
+            (d `shiftR` 1)
+            (min d (d + left))
+        pure (d', d' < d + left - 1)
+
+binsearch :: (Monad m) => (Int -> m Bool) -> Int -> Int -> m Int
+binsearch p = go
  where
-  go n m sq = do
-    let depth = m
-    let sq' = sq <> Seq.replicate depth True
-    let pth' = fromChoiceList (toList sq')
-    (i, l) <- lift_ (probeT rt pth')
-    if length l >= numberOfChoices pth' - depth + 1
-      then do
-        t <- p l i
-        if t
-          then go m (n + m) sq'
-          else case m of
-            1 -> go 1 1 (sq Seq.|> False)
-            m' -> go (m' - n) n sq
-      else pure i
-{-# INLINE reduceFibT #-}
+  go !m !h
+    | m >= pivot = do
+        pure m
+    | otherwise = do
+        p pivot >>= \case
+          True -> go pivot h
+          False -> go m pivot
+   where
+    pivot = (m + h) `shiftR` 1

rtree/src/Control/Monad/RTree.hs

@@ -33,12 +33,9 @@ module Control.Monad.RTree (
   module Data.RPath,
 ) where
 
-import Control.Applicative
 import Control.Monad
 import Control.Monad.Identity
 import Control.Monad.Reduce
-import Control.Monad.Trans
-import Control.Monad.Trans.Maybe (MaybeT (runMaybeT))
 import Data.Foldable
 import Data.Foldable.WithIndex
 import Data.Function ((&))
@@ -123,15 +120,15 @@ reduce
   :: (Monad m)
   => (i -> m Bool)
   -> RTree l i
-  -> m (Maybe i)
-reduce p = runMaybeT . checkgo
+  -> m i
+reduce p = go
  where
-  checkgo r = do
-    t <- lift $ p (extract r)
-    guard t *> go r
   go = \case
     Done i -> pure i
-    Split _ lhs rhs -> checkgo lhs <|> go rhs
+    Split _ lhs rhs -> do
+      p (extract lhs) >>= \case
+        True -> go lhs
+        False -> go rhs
 {-# INLINE reduce #-}
 
 -- | An RTreeT Node
@@ -180,14 +177,11 @@ reduceT
   -> (i -> n Bool)
   -> RTreeT l m i
   -> n i
-reduceT lift_ p = checkgo
- where
-  checkgo (RTreeT r) = do
-    r' <- lift_ r
-    t <- p =<< lift_ (extractN r')
-    unless t mzero
-    go r'
-  go = \case
+reduceT lift_ p = fix \rec x ->
+  lift_ (unRTreeT x) >>= \case
     DoneN i -> pure i
-    SplitN _ lhs rhs -> checkgo lhs <|> (lift_ (unRTreeT rhs) >>= go)
+    SplitN _ lhs rhs -> do
+      lift_ (extractT lhs) >>= p >>= \case
+        True -> rec lhs
+        False -> rec rhs
 {-# INLINE reduceT #-}

rtree/src/Data/RPath.hs

+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE LambdaCase #-}
+
 module Data.RPath (
   -- * RPath
   RPath,
@@ -6,6 +9,13 @@ module Data.RPath (
   numberOfChoices,
   indexChoice,
 
+  -- * Annotated
+  annotateChoices,
+  AnnotatedChoice (..),
+  Choice (..),
+  toDecision,
+  numberOfUndecided,
+
   -- * As a predicate
   toPredicate,
   toPredicateDebug,
@@ -16,6 +26,8 @@ module Data.RPath (
 ) where
 
 import Data.Bool
+import Data.Foldable
+import Data.Function
 import Data.IORef (atomicModifyIORef, newIORef)
 import Data.Maybe
 import Data.String
@@ -32,16 +44,16 @@ the rpath points to.
 -}
 toPredicate :: RPath -> IO (i -> IO Bool)
 toPredicate rp = do
-  idx <- newIORef (-1)
+  idx <- newIORef 0
   pure . const $ do
     idx' <- atomicModifyIORef idx (\n -> (n + 1, n))
     pure $ indexChoice rp idx'
 
 -- | Like @toPredicate@, but with debugging information
-toPredicateDebug :: (Show i) => RPath -> IO (i -> IO Bool)
+toPredicateDebug :: (Show i, Show l) => RPath -> IO ([AnnotatedChoice l] -> i -> IO Bool)
 toPredicateDebug rp = do
-  idx <- newIORef (-1)
-  pure $ \i -> do
+  idx <- newIORef 0
+  pure $ \_ i -> do
     idx' <- atomicModifyIORef idx (\n -> (n + 1, n))
     putStr (debugShowWithDepth rp idx')
     putStr ": "
@@ -52,9 +64,7 @@ toPredicateDebug rp = do
 default to false. Indexing with a negative number gives True.
 -}
 indexChoice :: RPath -> Int -> Bool
-indexChoice (RPath v) idx
-  | idx < 0 = True
-  | otherwise = fromMaybe False (v VU.!? idx)
+indexChoice (RPath v) idx = fromMaybe False (v VU.!? idx)
 {-# INLINE indexChoice #-}
 
 -- | Get the number of choices in the RPath.
@@ -77,6 +87,18 @@ debugShowWithDepth rp i =
   (map (bool '0' '1') . take (i + 1) . toChoiceList $ rp)
     <> replicate (numberOfChoices rp - i - 1) '*'
 
+-- | Get a list of annotated choices.
+annotateChoices :: RPath -> [l] -> [AnnotatedChoice l]
+annotateChoices rp =
+  zipWith
+    AnnotatedChoice
+    (map (bool No Yes) (toChoiceList rp) <> repeat Undecided)
+
+numberOfUndecided :: [AnnotatedChoice l] -> Int
+numberOfUndecided =
+  0 & foldl' \a ->
+    (\case Undecided -> 1 + a; _ow -> a) . choice
+
 debugShow :: RPath -> String
 debugShow = map (bool '0' '1') . toChoiceList
 
@@ -88,3 +110,20 @@ instance Read RPath where
 
 instance IsString RPath where
   fromString = fromChoiceList . map (== '1')
+
+-- | A Choice
+data Choice = Yes | No | Undecided
+  deriving (Show, Read, Eq, Ord, Enum, Bounded)
+
+-- | An annotated choice
+data AnnotatedChoice l = AnnotatedChoice
+  { choice :: !Choice
+  , label :: l
+  }
+  deriving (Show, Eq, Ord)
+
+toDecision :: Choice -> Maybe Bool
+toDecision = \case
+  Yes -> Just True
+  No -> Just False
+  Undecided -> Nothing

rtree/test/expected/double-let-expr-ired

+1***: y := 2; 1 + y True
+11***: 1 + 2 True
+111: 1 False
+1101: 2 False
+11001: 3 False
+11000: 1 + 2 True

rtree/test/expected/double-let-expr-ireda

+****: x := 1; y := 2; x + y True
+1***: y := 2; 1 + y True
+11***: 1 + 2 True
+111: 1 False

rtree/test/expected/double-overloading-let-expr-ired

+1***: x := 2; x + x True
+11***: 2 + 2 True
+111: 2 False
+1101: 2 False
+11001: 4 False
+11000: 2 + 2 True

rtree/test/expected/double-overloading-let-expr-ireda

+****: x := 1; x := 2; x + x True
+1***: x := 2; x + x True
+11***: 2 + 2 True
+111: 2 False

rtree/test/expected/small-let-expr-ired

+1***: 2 + 1 True
+11: 2 False
+101: 1 False
+1001: 3 False
+1000: 2 + 1 True

rtree/test/expected/small-let-expr-ireda

+***: x := 1; 2 + x True
+1***: 2 + 1 True
+11: 2 False

rtree/test/expected/small-opr-expr-ired

+1: 1 False
+01: 2 False
+001: 3 False
+000: 1 + 2 True

rtree/test/expected/small-opr-expr-ireda

+***: 1 + 2 True
+1: 1 False

rtree/test/src/Control/Monad/IRTreeSpec.hs

 {-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
 
 module Control.Monad.IRTreeSpec where
 
 import Control.Monad.IRTree
 import qualified Control.Monad.IRTree as IRTree
 import qualified Control.Monad.RTree as RTree
-import Control.Monad.Writer.Strict
 import Data.Bool
 import Data.Expr as Expr
 import Data.List.NonEmpty (nonEmpty)
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map.Strict as Map
+import System.IO
 import Test.Hspec
+import Test.Hspec.Glitter (onGlitter)
 
 spec :: Spec
-spec = describe "examples" do
+spec = do
+  binsearchSpec
+  expsearchSpec
+  examplesSpec
+
+binsearchSpec :: Spec
+binsearchSpec = describe "binsearch" do
+  it "should guess a number between 0 and 10" do
+    binsearch (\a -> do pure (a <= 6)) 0 10 `shouldReturn` 6
+    binsearch (\a -> do pure (a <= 3)) 0 10 `shouldReturn` 3
+
+expsearchSpec :: Spec
+expsearchSpec = describe "expsearch" do
+  it "should guess a number between 0 and 10" do
+    expsearch (\a -> do pure (a <= 6, 10 - a)) `shouldReturn` (6, True)
+    expsearch (\a -> do pure (a <= 3, 10 - a)) `shouldReturn` (3, True)
+    expsearch (\a -> do pure (a <= 9, 10 - a)) `shouldReturn` (9, False)
+    expsearch (\a -> do pure (a <= 12, 10 - a)) `shouldReturn` (9, False)
+
+examplesSpec :: Spec
+examplesSpec = describe "examples" do
   handle "small-opr-expr" $
     Opr (Cnt 1) (Cnt 2)
 
@@ -35,53 +58,85 @@ spec = describe "examples" do
 
     let re = runReaderT (Expr.rExpr e) Map.empty
 
+    it "should reduce like iinputs" do
+      forM_ (RTree.iinputs re) \(ii, e') -> do
+        p <- toPredicate ii
+        IRTree.reduce (const p) me `shouldReturn` e'
+
     let
       predicate :: Expr -> IO Bool
       predicate = pure . contains isOpr
 
     rex <- runIO $ RTree.reduce predicate re
 
-    -- onGlitterWith
-    --   ("test/expected/" <> str <> "-red")
-    --   ( \fp () -> do
-    --       (mex, result) <- runWriterT (IRTree.reduce (debugPredicate showString (prettyExprS 0) predicate) me)
-    --       writeFile fp (appEndo result "")
-    --       pure mex
-    --   )
-    --   do
-    --     it "should produce the same results as the RTree" \mex -> do
-    --       rex `shouldBe` mex
-
-    it "should find an opr exponentially" do
-      (mex, _) <- runWriterT (IRTree.reduceExp (debugPredicate showString (prettyExprS 0) predicate) me)
-      rex `shouldBe` mex
-
-    it "should find an opr fibonacci" do
-      (mex, _) <- runWriterT (IRTree.reduceFib (debugPredicate showString (prettyExprS 0) predicate) me)
-      rex `shouldBe` mex
-
-    it "should reduce like iinputs" do
-      forM_ (RTree.iinputs re) \(ii, e') -> do
-        p <- toPredicate ii
-        IRTree.reduce (const p) me `shouldReturn` Just e'
+    onGlitter
+      ("test/expected/" <> str <> "-ired")
+      ( \f -> do
+          withFile f WriteMode \h -> do
+            let p = debugPredicate' h predicate
+            IRTree.reduce p me `shouldReturn` rex
+      )
+      do
+        it "should find an opr normally" \_ -> do
+          IRTree.reduce (const predicate) me `shouldReturn` rex
+
+    onGlitter
+      ("test/expected/" <> str <> "-ireda")
+      ( \f -> do
+          withFile f WriteMode \h -> do
+            let p = debugPredicate' h predicate
+            let (i, l, _) = IRTree.probe me ""
+            _ <- p l i
+            IRTree.reduceAdaptive p me `shouldReturn` rex
+      )
+      do
+        it "should find an opr addabtively" \_ -> do
+          IRTree.reduceAdaptive (const predicate) me `shouldReturn` rex
+
+debugPredicate'
+  :: (Show l)
+  => Handle
+  -> (Expr -> IO Bool)
+  -> [AnnotatedChoice l]
+  -> Expr
+  -> IO Bool
+debugPredicate' h predicate lst i = do
+  x <- predicate i
+  hPutStr
+    h
+    ( map (debugChoice . choice) lst
+        <> ": "
+        <> prettyExprS 0 i ""
+        <> " "
+        <> show x
+        <> "\n"
+    )
+  pure x
+ where
+  debugChoice = \case
+    Yes -> '1'
+    No -> '0'
+    Undecided -> '*'
 
 debugPredicate
-  :: (Monad m)
-  => (l -> ShowS)
+  :: Handle
+  -> (l -> ShowS)
   -> (i -> ShowS)
-  -> (i -> m Bool)
-  -> [(Bool, l)]
+  -> (i -> IO Bool)
+  -> [AnnotatedChoice l]
   -> i
-  -> WriterT (Endo String) m Bool
-debugPredicate ppl ppi predicate lst i = do
-  x <- lift (predicate i)
-  tell . Endo $
-    showString (bool "0" "1" x)
-      . showString ": "
-      . ppi i
-      . showString "\n"
-      . case nonEmpty lst of
-        Nothing -> showString "initial\n"
-        Just lst' -> ppl (snd $ NE.last lst') . showString "\n"
-
+  -> IO Bool
+debugPredicate h ppl ppi predicate lst i = do
+  x <- predicate i
+  hPutStr
+    h
+    ( showString (bool "0" "1" x)
+        . showString ": "
+        . ppi i
+        . showString "\n"
+        . case nonEmpty lst of
+          Nothing -> showString "initial\n"
+          Just lst' -> ppl (label $ NE.last lst') . showString "\n"
+        $ ""
+    )
   pure x

rtree/test/src/Control/Monad/RTreeSpec.hs

@@ -115,4 +115,4 @@ examplesSpec = describe "example" do
     it "should reduce like iinputs" do
       forM_ (iinputs re) \(ii, e') -> do
         p <- toPredicate ii
-        reduce p re `shouldReturn` Just e'
+        reduce p re `shouldReturn` e'