EXPERIMENT: refactor graph traversal to replace leavesM

kore/kore.cabal

@@ -209,6 +209,7 @@ library
     ErrorContext
     From
     GlobalMain
+    GraphTraversal
     Injection
     Kore.AST.ApplicativeKore
     Kore.AST.AstWithLocation

kore/src/GraphTraversal.hs

@@ -0,0 +1,312 @@
+{-# LANGUAGE MultiWayIf #-}
+
+{- |
+Copyright   : (c) Runtime Verification, 2018-2022
+License     : BSD-3-Clause
+Maintainer  : jost.berthold@runtimeverification.com
+-}
+module GraphTraversal (
+    module GraphTraversal,
+) where
+
+import Control.Monad (foldM)
+import Control.Monad.Trans.State
+import Data.Limit
+import Data.List.NonEmpty qualified as NE
+import Data.Sequence (Seq (..))
+import Data.Sequence qualified as Seq
+import GHC.Generics qualified as GHC
+import GHC.Natural
+import Kore.Rewrite.Strategy (
+    GraphSearchOrder (..),
+    LimitExceeded (..),
+    TransitionT (..),
+    runTransitionT,
+    unfoldSearchOrder,
+ )
+import Kore.Simplify.Data (Simplifier)
+import Prelude.Kore
+import Pretty
+
+data TransitionResult a
+    = -- | straight-line execution
+      StraightLine a
+    | -- | branch point (1st arg), branching into 2nd arg elements
+      Branch a (NonEmpty a)
+    | -- | no next state but not final (e.g., not goal state, or side
+      -- conditions do not hold)
+      Stuck a
+    | -- | final state (e.g., goal state reached, side conditions hold)
+      Final a
+    | -- | not stuck, but also not final (maximum depth reached before
+      -- finishing the proof)
+      Stopped a
+    | -- Future work:
+
+      -- | config matches a terminal pattern (or: is RHS of a
+      -- "terminal" rule) Needs to return that RHS
+      Terminal a -- TODO(JB) obsolete, can use `Stopped`
+    | -- | config matches a cut pattern ( aka: is LHS of a "cut" rule)
+      -- The respective RHS (or RHSs) are returned (if any)
+      Cut a [a] -- TODO(JB) Could use `Stopped` if "next states" were added there
+    deriving stock (Eq, Show, GHC.Generic)
+
+instance Functor TransitionResult where
+    fmap f = \case
+        StraightLine a -> StraightLine $ f a
+        Branch a as -> Branch (f a) $ NE.map f as
+        Stuck a -> Stuck $ f a
+        Final a -> Final $ f a
+        Stopped a -> Stopped $ f a
+        Terminal a -> Terminal $ f a
+        Cut a as -> Cut (f a) (map f as)
+
+instance Pretty a => Pretty (TransitionResult a) where
+    pretty = \case
+        StraightLine a -> single "StraightLine" a
+        Branch a as -> multi "Branch" "node" a "successors" (NE.toList as)
+        Stuck a -> single "Stuck" a
+        Final a -> single "Final" a
+        Stopped a -> single "Stopped" a
+        Terminal a -> single "Terminal" a
+        Cut a as -> multi "Cut" "node" a "successors" as
+      where
+        single :: Doc x -> a -> Doc x
+        single lbl a =
+            Pretty.vsep [lbl, Pretty.indent 4 $ Pretty.pretty a]
+
+        multi :: Doc x -> Doc x -> a -> Doc x -> [a] -> Doc x
+        multi lbl lbl1 a lbl2 as =
+            Pretty.vsep $
+                [ lbl
+                , Pretty.indent 2 $ "- " <> lbl1
+                , Pretty.indent 4 $ Pretty.pretty a
+                , Pretty.indent 2 $ "- " <> lbl2
+                ]
+                    <> map (Pretty.indent 4 . Pretty.pretty) as
+
+-- Graph traversal would always stop at Terminal/Cut, and _may_ stop
+-- at Branch, depending on configuration.
+
+isStuck, isFinal, isStopped, isTerminal, isCut, isBranch :: TransitionResult a -> Bool
+isStuck (Stuck _) = True
+isStuck _ = False
+isFinal (Final _) = True
+isFinal _ = False
+isStopped (Stopped _) = True
+isStopped _ = False
+isBranch (Branch _ _) = True
+isBranch _ = False
+isTerminal (Terminal _) = True
+isTerminal _ = False
+isCut (Cut _ _) = True
+isCut _ = False
+
+extractNext :: TransitionResult a -> [a]
+extractNext = \case
+    StraightLine a -> [a]
+    Branch _ as -> NE.toList as
+    Stuck _ -> []
+    Final _ -> []
+    Stopped _ -> []
+    Terminal _ -> []
+    Cut _ as -> as
+
+extractState :: TransitionResult a -> Maybe a
+extractState = \case
+    StraightLine _ -> Nothing
+    Branch a _ -> Just a
+    Stuck a -> Just a
+    Final a -> Just a
+    Stopped a -> Just a
+    Terminal a -> Just a
+    Cut a _ -> Just a
+
+type Step prim = [prim]
+
+----------------------------------------
+transitionLeaves ::
+    forall c prim.
+    -- | Stop critera, in terms of 'TransitionResult's. The algorithm
+    -- will _always_ stop on 'Stuck' and 'Final', so [isTerminal,
+    -- isCut, isBranch] could be used here. Could simplify this to
+    -- FinalNodeType
+    [TransitionResult ([Step prim], c) -> Bool] ->
+    -- queue updating parameters,
+    -- we construct enqueue :: [a] -> Seq a -> m (Either LimitExceeded (Seq a)) from it
+    -- m is probably only there for logging purposes
+
+    -- | BreadthFirst, DepthFirst
+    GraphSearchOrder ->
+    -- | breadth limit, essentially a natural number
+    Limit Natural ->
+    -- | transition function
+    (([Step prim], c) -> Simplifier (TransitionResult ([Step prim], c))) ->
+    -- again, m is probably only for logging
+
+    -- | max-counterexamples, included in the internal logic
+    Natural ->
+    -- | initial node
+    ([Step prim], c) ->
+    Simplifier (TraversalResult c)
+transitionLeaves
+    stopCriteria
+    direction
+    breadthLimit
+    transit
+    maxCounterExamples
+    start =
+        enqueue [start] Seq.empty
+            >>= either
+                (pure . const (GotStuck 0 [Stopped $ snd start]))
+                (\q -> checkLeftUnproven <$> evalStateT (worker q) [])
+      where
+        enqueue' = unfoldSearchOrder direction
+
+        enqueue ::
+            [([Step prim], c)] ->
+            Seq ([Step prim], c) ->
+            Simplifier (Either (LimitExceeded ([Step prim], c)) (Seq ([Step prim], c)))
+        enqueue as q = do
+            newQ <- enqueue' as q
+            pure $
+                if exceedsLimit newQ
+                    then Left (LimitExceeded newQ)
+                    else Right newQ
+
+        exceedsLimit :: Seq ([Step prim], c) -> Bool
+        exceedsLimit = not . withinLimit breadthLimit . fromIntegral . Seq.length
+
+        shouldStop :: TransitionResult ([Step prim], c) -> Bool
+        shouldStop result = any ($ result) stopCriteria
+
+        maxStuck = fromIntegral maxCounterExamples
+
+        worker ::
+            Seq ([Step prim], c) ->
+            StateT
+                [TransitionResult ([Step prim], c)]
+                Simplifier
+                (TraversalResult ([Step prim], c))
+        worker Seq.Empty = Ended . reverse <$> get
+        worker (a :<| as) = do
+            result <- lift $ step a as
+            case result of
+                Continue nextQ -> worker nextQ
+                Output oneResult nextQ -> do
+                    modify (oneResult :)
+                    if not (isStuck oneResult)
+                        then worker nextQ
+                        else do
+                            stuck <- gets (filter isStuck)
+                            if length stuck >= maxStuck
+                                then
+                                    pure $
+                                        GotStuck (Seq.length nextQ) stuck
+                                else worker nextQ
+                Abort lastState queue -> do
+                    pure $ Aborted (Seq.length queue) [lastState]
+        -- TODO could add current state to return value ^^^^^^^
+        step ::
+            ([Step prim], c) ->
+            Seq ([Step prim], c) ->
+            Simplifier (StepResult ([Step prim], c))
+        step a q = do
+            next <- transit a
+            if (isStuck next || isFinal next || isStopped next || shouldStop next)
+                then pure (Output next q)
+                else
+                    let abort (LimitExceeded queue) = Abort next queue
+                     in either abort Continue <$> enqueue (extractNext next) q
+
+        checkLeftUnproven ::
+            TraversalResult ([Step prim], c) -> TraversalResult c
+        checkLeftUnproven = \case
+            result@(Ended results) ->
+                let -- we collect a maximum of 'maxCounterExamples' Stuck states
+                    stuck = map (fmap snd) $ filter isStuck results
+                    -- Other states may be unfinished but not stuck (Stopped)
+                    -- Only provide the requested amount of states (maxCounterExamples)
+                    unproven = take maxStuck . map (fmap snd) $ filter isStopped results
+                 in if
+                            | (not $ null stuck) -> GotStuck 0 stuck
+                            | not $ null unproven -> GotStuck 0 unproven
+                            | otherwise -> fmap snd result
+            other -> fmap snd other
+
+data StepResult a
+    = -- | Traversal continues with given queue.
+      Continue (Seq a)
+    | -- | Traversal produced a result and continues with given queue.
+      -- Typically a final or stuck state (or a "stop" state), and the
+      -- queue is the remaining work.
+      Output (TransitionResult a) (Seq a)
+    | -- | Traversal exceeded the breadth limit and should not
+      -- continue. The last state and the current queue are provided
+      -- for analysis.
+      Abort (TransitionResult a) (Seq a)
+    deriving stock (Eq, Show, GHC.Generic)
+
+-- TODO(JB) extract and return states instead of TransitionResults
+data TraversalResult a
+    = -- | remaining queue length and stuck or stopped (unproven)
+      -- results (always at most maxCounterExamples many). Caller
+      -- should extract current states from returned results.
+      GotStuck Int [TransitionResult a]
+    | -- | queue length (exceeding the limit) and result(s) of the
+      -- last step that led to stopping. Caller should extract
+      -- following states from the result.
+      Aborted Int [TransitionResult a]
+    | -- | queue empty, results returned
+      Ended [TransitionResult a]
+    deriving stock (Eq, Show, GHC.Generic)
+
+instance Pretty a => Pretty (TraversalResult a) where
+    pretty = \case
+        GotStuck n as ->
+            Pretty.hang 4 . Pretty.vsep $
+                ("Got stuck with queue of " <> Pretty.pretty n) :
+                map Pretty.pretty as
+        Aborted n as ->
+            Pretty.hang 4 . Pretty.vsep $
+                ("Aborted with queue of " <> Pretty.pretty n) :
+                map Pretty.pretty as
+        Ended as ->
+            Pretty.hang 4 . Pretty.vsep $
+                "Ended" : map Pretty.pretty as
+
+instance Functor TraversalResult where
+    fmap f = \case
+        GotStuck n rs -> GotStuck n (ffmap f rs)
+        Aborted n rs -> Aborted n (ffmap f rs)
+        Ended rs -> Ended (ffmap f rs)
+      where
+        ffmap = map . fmap
+
+----------------------------------------
+-- constructing transition functions (for caller)
+
+simpleTransition ::
+    forall config m prim rule.
+    Monad m =>
+    -- | primitive strategy rule
+    (prim -> config -> TransitionT rule m config) ->
+    -- | converter to interpret the config (claim state or program state)
+    (config -> [config] -> TransitionResult config) ->
+    -- TODO(JB) should also consider the applied rule(s) (for Terminal/Cut)
+    -- final transition function
+    ([Step prim], config) ->
+    m (TransitionResult ([Step prim], config))
+simpleTransition applyPrim mapToResult = uncurry tt
+  where
+    tt :: [Step prim] -> config -> m (TransitionResult ([Step prim], config))
+    tt [] config =
+        pure $ fmap ([],) $ mapToResult config []
+    tt ([] : iss) config =
+        tt iss config
+    tt (is : iss) config =
+        (fmap (iss,) . mapToResult config . map fst)
+            <$> runTransitionT (applyGroup is config)
+
+    applyGroup :: Step prim -> config -> TransitionT rule m config
+    applyGroup is c = foldM (flip applyPrim) c is