diff --git a/pyrtree/rtree.py b/pyrtree/rtree.py index bec022bde18d0a6af2ef8d2610f77ddcf4c93eaf..d484af321b37873da6e0b2ce55245346dab7defc 100644 --- a/pyrtree/rtree.py +++ b/pyrtree/rtree.py @@ -11,8 +11,8 @@ class ReducePath: self.path.append(choice) return self - def didGuess(self): - return self.index > len(self.path) + def left(self): + return self.index - len(self.path) def dispensable(self): self.index += 1 @@ -25,16 +25,16 @@ class ReducePath: def reduce(predicate, rtree): r = ReducePath([]) i = rtree(r) - - if not predicate(i): - return None - - while r.didGuess(): - # Explore the left tree - i = rtree(r.explore(True)) - # If the predcate fails, move right - r.path[-1] = predicate(i) - + it = rtree(r.explore(True)) + # While we don't consume all choices going down the true branch + while r.left() >= 0: + if predicate(it): + # If true update the valid input + i = it + else: + # If false we have to go down the left branch. + r.path[-1] = False + it = rtree(r.explore(True)) return i diff --git a/rtree/src/Control/Monad/IRTree.hs b/rtree/src/Control/Monad/IRTree.hs index 706c9dbd7db92a9b175d23cc960b1cc20cf3c40c..b8c3a22bce1ef272af7ff564bbc1001f3d8dd748 100644 --- a/rtree/src/Control/Monad/IRTree.hs +++ b/rtree/src/Control/Monad/IRTree.hs @@ -92,15 +92,14 @@ reduceT -> IRTreeT l t i -> m i reduceT lift_ p rt = do - Seq.empty & fix \rec sq -> do - -- Try to run the true branch. + (k', _, _) <- _probe Seq.empty + (\f -> f Seq.empty k') $ fix \rec sq k -> do (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 + if left < 0 + then pure k + else do + t <- p l i + rec (sq Seq.|> t) (if t then i else k) where _probe sq = lift_ . probeT rt . fromChoiceList $ toList sq {-# INLINE reduceT #-} diff --git a/rtree/test/expected/double-let-expr-ired b/rtree/test/expected/double-let-expr-ired index 3dce08b5a6bd98c8c1e2023c4d7195685f12f41d..35bdda26d87fbec65f9fb9b1cad67cb82125c8b8 100644 --- a/rtree/test/expected/double-let-expr-ired +++ b/rtree/test/expected/double-let-expr-ired @@ -3,4 +3,3 @@ 111: 1 False 1101: 2 False 11001: 3 False -11000: 1 + 2 True diff --git a/rtree/test/expected/double-overloading-let-expr-ired b/rtree/test/expected/double-overloading-let-expr-ired index 5d756a44ef91e13e43a83c809f836551fb4d4431..4f29850714ad4f0fd7069a26d78919b5d0fb84dc 100644 --- a/rtree/test/expected/double-overloading-let-expr-ired +++ b/rtree/test/expected/double-overloading-let-expr-ired @@ -3,4 +3,3 @@ 111: 2 False 1101: 2 False 11001: 4 False -11000: 2 + 2 True diff --git a/rtree/test/expected/small-let-expr-ired b/rtree/test/expected/small-let-expr-ired index 8ec233528b56f31620c43cab29f50b2a5eb1f9e2..ab4bde13d3d2bdb77639c66c5e4225542e53d139 100644 --- a/rtree/test/expected/small-let-expr-ired +++ b/rtree/test/expected/small-let-expr-ired @@ -2,4 +2,3 @@ 11: 2 False 101: 1 False 1001: 3 False -1000: 2 + 1 True diff --git a/rtree/test/expected/small-opr-expr-ired b/rtree/test/expected/small-opr-expr-ired index e8e5ee3c331fe0a608e4f23095e76af30269bc65..63b780a5796ded28dd8d789958a18af8241015a8 100644 --- a/rtree/test/expected/small-opr-expr-ired +++ b/rtree/test/expected/small-opr-expr-ired @@ -1,4 +1,3 @@ 1: 1 False 01: 2 False 001: 3 False -000: 1 + 2 True