From 4bcb45c077fdb4201aa6a7c3ced54a17acfd5971 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Tue, 16 Jul 2024 13:17:55 -0700
Subject: [PATCH 1/6] add chooseInput to TraceTree to provide arbitrary input

---
 src/Pate/TraceTree.hs | 55 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/src/Pate/TraceTree.hs b/src/Pate/TraceTree.hs
index 4fdc60da..3e3f25b9 100644
--- a/src/Pate/TraceTree.hs
+++ b/src/Pate/TraceTree.hs
@@ -34,6 +34,7 @@ type class with a (distinct) symbol.
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE QuantifiedConstraints #-}
 {-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE TypeOperators #-}
 
 module Pate.TraceTree (
     TraceTree
@@ -112,6 +113,9 @@ module Pate.TraceTree (
   , forkTraceTreeHook
   , NodeFinalAction(..)
   , QueryResult(..)
+  , InputChoice
+  , giveChoiceInput
+  , chooseInput
   ) where
 
 import           GHC.TypeLits ( Symbol, KnownSymbol )
@@ -1007,6 +1011,57 @@ chooseLazy treenm f = do
           False -> return Nothing
     DefaultChoice -> return $ LazyIOAction (return False) (\_ -> return Nothing)
 
+data InputChoice (k :: l) nm where
+  InputChoice :: (IsTraceNode k nm, TraceNodeLabel nm ~ ()) =>
+    { inputChoiceParse :: String -> Maybe (TraceNodeType k nm) -- FIXME: make this part of the type class?
+    , inputChoicePut :: (TraceNodeType k nm) -> IO Bool -- returns false if input has already been provided
+    , inputChoiceValue :: IO (Maybe (TraceNodeType k nm))
+    } -> InputChoice k nm
+
+giveChoiceInput :: InputChoice k nm -> String -> IO Bool
+giveChoiceInput ic input = inputChoiceValue ic >>= \case
+  Just{} -> return False
+  Nothing -> case inputChoiceParse ic input of
+    Just v -> inputChoicePut ic v
+    Nothing -> return False
+
+instance IsTraceNode k "choiceInput" where
+  type TraceNodeType k "choiceInput" = Some (InputChoice k)
+  type TraceNodeLabel "choiceInput" = String
+  prettyNode lbl _ = PP.pretty lbl
+  nodeTags = mkTags @k @"choiceInput" [Summary, Simplified]
+  jsonNode core lbl (Some (ic@InputChoice{} :: InputChoice k nm)) = do
+    v_json <- inputChoiceValue ic >>= \case
+      Just v -> jsonNode @_ @k @nm core () v
+      Nothing -> return JSON.Null
+    return $ JSON.object ["node_kind" JSON..= ("choiceInput" :: String), "value" JSON..= v_json, "prompt" JSON..= lbl]
+
+chooseInput ::
+  forall nm_choice a b k m e.
+  IsTreeBuilder k e m =>
+  IsTraceNode k nm_choice =>
+  IO.MonadUnliftIO m =>
+  TraceNodeLabel nm_choice ~ () =>
+  String ->
+  (String -> Maybe (TraceNodeType k nm_choice)) ->
+  (TraceNodeType k nm_choice -> b -> IO a) ->
+  m (LazyIOAction b a)
+chooseInput treenm parseInput f = do
+  builder <- getTreeBuilder
+  case interactionMode builder of
+    Interactive{} -> do
+      c <- liftIO $ newEmptyMVar
+      let getValue = tryReadMVar c
+      let putValue v = tryPutMVar c v
+      let isReady = not <$> isEmptyMVar c
+      let ichoice = InputChoice @k @nm_choice parseInput putValue getValue
+      emitTraceLabel @"choiceInput" @k treenm (Some ichoice)
+      return $ LazyIOAction (liftIO isReady) $ \inputVal -> getValue >>= \case
+        Just v -> Just <$> f v inputVal
+        Nothing -> return Nothing
+    DefaultChoice{} -> return $ LazyIOAction (return False) (\_ -> return Nothing)
+         
+
 choose ::
   forall nm_choice a k m e.
   IsTreeBuilder k e m =>

From 3b5a0964df17eaac4ca16e1d5ad14c0c7a1f3b05 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Wed, 17 Jul 2024 10:07:20 -0700
Subject: [PATCH 2/6] TraceTree: allow input nodes to have any label type

---
 src/Pate/TraceTree.hs | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/src/Pate/TraceTree.hs b/src/Pate/TraceTree.hs
index 3e3f25b9..208b73ce 100644
--- a/src/Pate/TraceTree.hs
+++ b/src/Pate/TraceTree.hs
@@ -1012,17 +1012,17 @@ chooseLazy treenm f = do
     DefaultChoice -> return $ LazyIOAction (return False) (\_ -> return Nothing)
 
 data InputChoice (k :: l) nm where
-  InputChoice :: (IsTraceNode k nm, TraceNodeLabel nm ~ ()) =>
-    { inputChoiceParse :: String -> Maybe (TraceNodeType k nm) -- FIXME: make this part of the type class?
-    , inputChoicePut :: (TraceNodeType k nm) -> IO Bool -- returns false if input has already been provided
-    , inputChoiceValue :: IO (Maybe (TraceNodeType k nm))
+  InputChoice :: (IsTraceNode k nm) =>
+    { inputChoiceParse :: String -> Maybe (TraceNodeLabel nm, TraceNodeType k nm) -- FIXME: make this part of the type class?
+    , inputChoicePut :: TraceNodeLabel nm -> TraceNodeType k nm -> IO Bool -- returns false if input has already been provided
+    , inputChoiceValue :: IO (Maybe (TraceNodeLabel nm, TraceNodeType k nm))
     } -> InputChoice k nm
 
 giveChoiceInput :: InputChoice k nm -> String -> IO Bool
 giveChoiceInput ic input = inputChoiceValue ic >>= \case
   Just{} -> return False
   Nothing -> case inputChoiceParse ic input of
-    Just v -> inputChoicePut ic v
+    Just (lbl, v) -> inputChoicePut ic lbl v
     Nothing -> return False
 
 instance IsTraceNode k "choiceInput" where
@@ -1032,7 +1032,7 @@ instance IsTraceNode k "choiceInput" where
   nodeTags = mkTags @k @"choiceInput" [Summary, Simplified]
   jsonNode core lbl (Some (ic@InputChoice{} :: InputChoice k nm)) = do
     v_json <- inputChoiceValue ic >>= \case
-      Just v -> jsonNode @_ @k @nm core () v
+      Just (lbl', v) -> jsonNode @_ @k @nm core lbl' v
       Nothing -> return JSON.Null
     return $ JSON.object ["node_kind" JSON..= ("choiceInput" :: String), "value" JSON..= v_json, "prompt" JSON..= lbl]
 
@@ -1041,10 +1041,9 @@ chooseInput ::
   IsTreeBuilder k e m =>
   IsTraceNode k nm_choice =>
   IO.MonadUnliftIO m =>
-  TraceNodeLabel nm_choice ~ () =>
   String ->
-  (String -> Maybe (TraceNodeType k nm_choice)) ->
-  (TraceNodeType k nm_choice -> b -> IO a) ->
+  (String -> Maybe (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
+  (TraceNodeLabel nm_choice -> TraceNodeType k nm_choice -> b -> IO a) ->
   m (LazyIOAction b a)
 chooseInput treenm parseInput f = do
   builder <- getTreeBuilder
@@ -1052,15 +1051,14 @@ chooseInput treenm parseInput f = do
     Interactive{} -> do
       c <- liftIO $ newEmptyMVar
       let getValue = tryReadMVar c
-      let putValue v = tryPutMVar c v
+      let putValue lbl v = tryPutMVar c (lbl, v)
       let isReady = not <$> isEmptyMVar c
       let ichoice = InputChoice @k @nm_choice parseInput putValue getValue
       emitTraceLabel @"choiceInput" @k treenm (Some ichoice)
       return $ LazyIOAction (liftIO isReady) $ \inputVal -> getValue >>= \case
-        Just v -> Just <$> f v inputVal
+        Just (lbl, v) -> Just <$> f lbl v inputVal
         Nothing -> return Nothing
     DefaultChoice{} -> return $ LazyIOAction (return False) (\_ -> return Nothing)
-         
 
 choose ::
   forall nm_choice a k m e.

From 0a9b4580414ec7856cb0c61d4c51942bf1eee203 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Wed, 17 Jul 2024 10:12:17 -0700
Subject: [PATCH 3/6] TraceTree: add error reporting to choiceInput

---
 src/Pate/TraceTree.hs | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/Pate/TraceTree.hs b/src/Pate/TraceTree.hs
index 208b73ce..b796d9e9 100644
--- a/src/Pate/TraceTree.hs
+++ b/src/Pate/TraceTree.hs
@@ -1011,19 +1011,26 @@ chooseLazy treenm f = do
           False -> return Nothing
     DefaultChoice -> return $ LazyIOAction (return False) (\_ -> return Nothing)
 
+data InputChoiceError = 
+    InputChoiceAlreadyMade
+  | InputChoiceParseError String
+  deriving Show
+
 data InputChoice (k :: l) nm where
   InputChoice :: (IsTraceNode k nm) =>
-    { inputChoiceParse :: String -> Maybe (TraceNodeLabel nm, TraceNodeType k nm) -- FIXME: make this part of the type class?
+    { inputChoiceParse :: String -> Either String (TraceNodeLabel nm, TraceNodeType k nm)
     , inputChoicePut :: TraceNodeLabel nm -> TraceNodeType k nm -> IO Bool -- returns false if input has already been provided
     , inputChoiceValue :: IO (Maybe (TraceNodeLabel nm, TraceNodeType k nm))
     } -> InputChoice k nm
 
-giveChoiceInput :: InputChoice k nm -> String -> IO Bool
+giveChoiceInput :: InputChoice k nm -> String -> IO (Maybe InputChoiceError)
 giveChoiceInput ic input = inputChoiceValue ic >>= \case
-  Just{} -> return False
+  Just{} -> return $ Just InputChoiceAlreadyMade
   Nothing -> case inputChoiceParse ic input of
-    Just (lbl, v) -> inputChoicePut ic lbl v
-    Nothing -> return False
+    Right (lbl, v) -> inputChoicePut ic lbl v >>= \case
+      True -> return Nothing
+      False -> return $ Just InputChoiceAlreadyMade 
+    Left err -> return $ Just $ InputChoiceParseError err
 
 instance IsTraceNode k "choiceInput" where
   type TraceNodeType k "choiceInput" = Some (InputChoice k)
@@ -1042,7 +1049,7 @@ chooseInput ::
   IsTraceNode k nm_choice =>
   IO.MonadUnliftIO m =>
   String ->
-  (String -> Maybe (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
+  (String -> Either String (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
   (TraceNodeLabel nm_choice -> TraceNodeType k nm_choice -> b -> IO a) ->
   m (LazyIOAction b a)
 chooseInput treenm parseInput f = do

From 068c4f1cee5e95e70249751d1e77b43d3ae16984 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Mon, 22 Jul 2024 12:22:52 -0700
Subject: [PATCH 4/6] add 'input' command to Repl for supplying string input

---
 src/Pate/TraceTree.hs                   | 89 ++++++++++++++++++++++---
 src/Pate/Verification/StrongestPosts.hs | 27 ++++----
 tools/pate-repl/Repl.hs                 | 26 ++++++++
 3 files changed, 120 insertions(+), 22 deletions(-)

diff --git a/src/Pate/TraceTree.hs b/src/Pate/TraceTree.hs
index b796d9e9..182909aa 100644
--- a/src/Pate/TraceTree.hs
+++ b/src/Pate/TraceTree.hs
@@ -114,8 +114,12 @@ module Pate.TraceTree (
   , NodeFinalAction(..)
   , QueryResult(..)
   , InputChoice
+  , InputChoiceError(InputChoiceError)
   , giveChoiceInput
+  , waitingForChoiceInput
   , chooseInput
+  , chooseInputLazy
+  , chooseInputFromList
   ) where
 
 import           GHC.TypeLits ( Symbol, KnownSymbol )
@@ -128,7 +132,7 @@ import           Data.String
 import qualified Data.Map as Map
 import           Data.Map ( Map )
 import           Data.Default
-import           Data.List ((!!), find, isPrefixOf)
+import           Data.List (isPrefixOf, findIndex)
 import           Control.Monad.Trans (lift)
 import           Control.Monad.Trans.Maybe
 import qualified Control.Monad.Reader as CMR
@@ -1013,24 +1017,36 @@ chooseLazy treenm f = do
 
 data InputChoiceError = 
     InputChoiceAlreadyMade
-  | InputChoiceParseError String
+  | InputChoiceError String [String]
   deriving Show
 
+
+instance PP.Pretty InputChoiceError where
+  pretty = \case
+    InputChoiceAlreadyMade -> "INTERNAL ERROR: input choice has already been given!"
+    InputChoiceError msgHeader [] -> PP.pretty msgHeader
+    InputChoiceError msgHeader msgDetails -> PP.pretty msgHeader PP.<> PP.line PP.<> PP.vsep (map PP.pretty msgDetails)  
+
 data InputChoice (k :: l) nm where
   InputChoice :: (IsTraceNode k nm) =>
-    { inputChoiceParse :: String -> Either String (TraceNodeLabel nm, TraceNodeType k nm)
+    { inputChoiceParse :: String -> Either InputChoiceError (TraceNodeLabel nm, TraceNodeType k nm)
     , inputChoicePut :: TraceNodeLabel nm -> TraceNodeType k nm -> IO Bool -- returns false if input has already been provided
     , inputChoiceValue :: IO (Maybe (TraceNodeLabel nm, TraceNodeType k nm))
     } -> InputChoice k nm
 
+waitingForChoiceInput :: InputChoice k nm -> IO Bool
+waitingForChoiceInput ic = inputChoiceValue ic >>= \case
+  Just{} -> return False
+  Nothing -> return True
+
 giveChoiceInput :: InputChoice k nm -> String -> IO (Maybe InputChoiceError)
-giveChoiceInput ic input = inputChoiceValue ic >>= \case
-  Just{} -> return $ Just InputChoiceAlreadyMade
-  Nothing -> case inputChoiceParse ic input of
+giveChoiceInput ic input = waitingForChoiceInput ic >>= \case
+  False -> return $ Just InputChoiceAlreadyMade
+  True -> case inputChoiceParse ic input of
     Right (lbl, v) -> inputChoicePut ic lbl v >>= \case
       True -> return Nothing
       False -> return $ Just InputChoiceAlreadyMade 
-    Left err -> return $ Just $ InputChoiceParseError err
+    Left err -> return $ Just err
 
 instance IsTraceNode k "choiceInput" where
   type TraceNodeType k "choiceInput" = Some (InputChoice k)
@@ -1043,16 +1059,71 @@ instance IsTraceNode k "choiceInput" where
       Nothing -> return JSON.Null
     return $ JSON.object ["node_kind" JSON..= ("choiceInput" :: String), "value" JSON..= v_json, "prompt" JSON..= lbl]
 
+instance IsTraceNode k "opt_index" where
+  type TraceNodeType k "opt_index" = Int
+  type TraceNodeLabel "opt_index" = String
+  prettyNode msg _ = PP.pretty msg
+  nodeTags = (mkTags @k @"opt_index" [Summary, Simplified]) ++ 
+    [("debug", \msg lbl -> PP.pretty lbl PP.<> ":" PP.<+> PP.pretty msg)]
+
+-- | Wrapper for 'chooseInput' that just takes a list of labeled options
+--   and generates a parser for picking one option.
+chooseInputFromList ::
+  IsTreeBuilder k e m =>
+  IO.MonadUnliftIO m =>
+  String ->
+  [(String, a)] ->
+  m (Maybe a)
+chooseInputFromList treenm opts = do
+  let parseInput s = case findIndex (\(s',_) -> s == s') opts of
+        Just idx -> Right (s, idx)
+        Nothing -> Left (InputChoiceError "Invalid input. Valid options:" (map fst opts))
+  chooseInput @"opt_index" treenm parseInput >>= \case
+    Just (_, idx) -> return $ Just $ (snd (opts !! idx))
+    Nothing -> return Nothing
+
+-- | Take user input as a string. Returns 'Nothing' in the case where the trace tree
+--   is not running interactively. Otherwise, blocks the current thread until
+--   valid input is provided.
 chooseInput ::
+  forall nm_choice k m e.
+  IsTreeBuilder k e m =>
+  IsTraceNode k nm_choice =>
+  IO.MonadUnliftIO m =>
+  String ->
+  (String -> Either InputChoiceError (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
+  m (Maybe (TraceNodeLabel nm_choice, TraceNodeType k nm_choice))
+chooseInput treenm parseInput = do
+  builder <- getTreeBuilder
+  case interactionMode builder of
+    Interactive nextChoiceIdent -> do
+      newChoiceIdent <- liftIO $ nextChoiceIdent
+      let status = NodeStatus StatusSuccess False (BlockedStatus (Set.singleton newChoiceIdent) Set.empty)
+      let statusFinal = NodeStatus StatusSuccess True (BlockedStatus Set.empty (Set.singleton newChoiceIdent))
+      c <- liftIO $ newEmptyMVar
+      let getValue = tryReadMVar c
+      let putValue lbl v = tryPutMVar c (lbl, v)
+      let ichoice = InputChoice @k @nm_choice parseInput putValue getValue
+      (lbl, v) <- withTracingLabel @"choiceInput" @k treenm (Some ichoice) $ do
+        builder' <- getTreeBuilder
+        liftIO $ updateTreeStatus builder' status
+        (lbl, v) <- liftIO $ readMVar c
+        emitTraceLabel @nm_choice lbl v
+        liftIO $ updateTreeStatus builder' statusFinal
+        return (lbl, v)
+      return $ Just (lbl, v)
+    DefaultChoice -> return Nothing
+
+chooseInputLazy ::
   forall nm_choice a b k m e.
   IsTreeBuilder k e m =>
   IsTraceNode k nm_choice =>
   IO.MonadUnliftIO m =>
   String ->
-  (String -> Either String (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
+  (String -> Either InputChoiceError (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
   (TraceNodeLabel nm_choice -> TraceNodeType k nm_choice -> b -> IO a) ->
   m (LazyIOAction b a)
-chooseInput treenm parseInput f = do
+chooseInputLazy treenm parseInput f = do
   builder <- getTreeBuilder
   case interactionMode builder of
     Interactive{} -> do
diff --git a/src/Pate/Verification/StrongestPosts.hs b/src/Pate/Verification/StrongestPosts.hs
index daa3ef3a..62b45e03 100644
--- a/src/Pate/Verification/StrongestPosts.hs
+++ b/src/Pate/Verification/StrongestPosts.hs
@@ -500,25 +500,26 @@ pickCutPoints pickMany msg inputs = go []
     hasBin bin picked = 
       any (\(Some (PPa.WithBin bin' _)) -> case testEquality bin bin' of Just Refl -> True; _ -> False) picked
     go picked = do
-      mres <- choose @"()" msg $ \choice -> do
-        forM_ inputs $ \(_divergeSingle, Some blk, PD.ParsedBlocks pblks) -> do
-          let bin = PB.blockBinRepr blk
-          forM_ pblks $ \pblk -> forM_ (getIntermediateAddrs pblk) $ \addr -> do
-            -- FIXME: block entry kind is unused at the moment?
-            let concBlk = PB.mkConcreteBlock blk PB.BlockEntryJump addr
-            --let node = mkNodeEntry' divergeSingle (PPa.mkSingle bin concBlk)
-            choice (show addr ++ " " ++ "(" ++ show bin ++ ")") () $ do
-              return $ Just $ (Pair concBlk (PPa.WithBin bin addr))
-        case pickMany && hasBin PBi.OriginalRepr picked && hasBin PBi.PatchedRepr picked of
-          True -> choice "Finish Choosing" () $ return Nothing 
-          False -> return ()
+      let addr_opts = 
+           [ (show addr ++ " " ++ "(" ++ show bin ++ ")", Just $ Pair concBlk (PPa.WithBin bin addr))
+           | (_divergeSingle, Some blk, PD.ParsedBlocks pblks) <- inputs
+           , bin <- return $ PB.blockBinRepr blk
+           , pblk <- pblks
+           , addr <- getIntermediateAddrs pblk
+           , concBlk <- return $ PB.mkConcreteBlock blk PB.BlockEntryJump addr
+           ]
+      let opts = case pickMany && hasBin PBi.OriginalRepr picked && hasBin PBi.PatchedRepr picked of
+            True -> addr_opts ++ [("Finished Choosing", Nothing), ("", Nothing)]
+            False -> addr_opts
+      mres <- chooseInputFromList msg opts
       case mres of
-        Just (Pair blk (PPa.WithBin bin addr)) -> do
+        Just (Just (Pair blk (PPa.WithBin bin addr))) -> do
           _ <- addIntraBlockCut addr blk
           let x = Some (PPa.WithBin bin (PAd.segOffToAddr addr))
           case pickMany of
             True -> go (x:picked)
             False -> return (x:picked)
+        Just Nothing -> return picked
         Nothing -> return picked
 
 
diff --git a/tools/pate-repl/Repl.hs b/tools/pate-repl/Repl.hs
index 0f820be6..12cc4ba3 100644
--- a/tools/pate-repl/Repl.hs
+++ b/tools/pate-repl/Repl.hs
@@ -732,6 +732,32 @@ goto idx = execReplM $ do
     Just _ -> return ()
     Nothing -> PO.printErrLn "No such option"
 
+asInput :: forall sym arch nm. TraceNode sym arch nm -> ReplM sym arch (Maybe (Some (InputChoice '(sym,arch))))
+asInput (node@(TraceNode _ v _)) = case testEquality (knownSymbol @nm) (knownSymbol @"choiceInput")  of
+  Just Refl -> do
+    Some tr <- return $ v
+    (IO.liftIO $ waitingForChoiceInput tr) >>= \case
+      True -> return $ Just v
+      False -> return Nothing
+  Nothing -> return Nothing
+
+input' :: String -> ReplM sym arch (Maybe InputChoiceError)
+input' s = withCurrentNode $ \tr -> asInput tr >>= \case
+  Just (Some ic) -> (IO.liftIO $ giveChoiceInput ic s) >>= \case
+    Nothing -> do
+      IO.liftIO $ IO.threadDelay 100
+      top'
+      IO.liftIO $ wait
+      return Nothing
+    Just err -> return $ Just err
+  Nothing -> return $ Just (InputChoiceError "Not an input prompt" [])
+
+input :: String -> IO ()
+input s = execReplM $ do
+  input' s >>= \case
+    Just err -> PO.printErrLn (PP.pretty err)
+    Nothing -> return ()
+
 gotoIndex :: forall sym arch. Integer -> ReplM sym arch String
 gotoIndex idx = (goto' (fromIntegral idx)) >>= \case
   Just (Some ((TraceNode lbl v _) :: TraceNode sym arch nm)) -> do

From e752d06e1133ca08980ee419b73cf28cfa71df64 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Wed, 24 Jul 2024 13:44:50 -0700
Subject: [PATCH 5/6] support providing input from scripts

---
 src/Pate/Script.hs                      | 28 ++++++--
 src/Pate/TraceTree.hs                   | 96 +++++++++++++++----------
 src/Pate/Verification/StrongestPosts.hs |  2 +-
 3 files changed, 82 insertions(+), 44 deletions(-)

diff --git a/src/Pate/Script.hs b/src/Pate/Script.hs
index 27cce9b4..4a0c1d15 100644
--- a/src/Pate/Script.hs
+++ b/src/Pate/Script.hs
@@ -68,6 +68,7 @@ the rest of the query.
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE MultiWayIf #-}
 
 module Pate.Script (
     readScript
@@ -91,6 +92,9 @@ import           Text.Megaparsec ((<|>))
 import           Pate.TraceTree
 import Data.Void
 import Data.Maybe (fromMaybe)
+import Data.Parameterized (Some(..))
+import qualified System.IO as IO
+import qualified Prettyprinter as PP
 -- import qualified System.IO as IO
 
 data Script = Script [NodeQuery]
@@ -127,16 +131,30 @@ data ScriptStep =
  | ScriptTerminal NodeIdentQuery NodeFinalAction
 
 pickChoiceAct :: NodeFinalAction
-pickChoiceAct = NodeFinalAction $ \node -> case asChoice node of
-  Just (SomeChoice c) -> choiceChosen c >>= \case
+pickChoiceAct = NodeFinalAction $ \remaining node -> if
+  | Just (SomeChoice c) <- asChoice node
+  , [] <- remaining
+  -> choiceChosen c >>= \case
     True -> return False
     False -> choiceReady (choiceHeader c) >>= \case
       True -> return False
       False -> choicePick c >> return True
-  Nothing -> return False
-
+  | Just (Some ic) <- asInputChoice node
+    -- if we're at a choiceInput node and have exactly one string query remaining,
+    -- then we attempt to give this as input and match if it is accepted by the parser
+  , [QueryString s] <- remaining
+  -> giveChoiceInput ic s >>= \case
+    -- no error means the input was accepted, which indicates a successful match
+    Nothing -> return True
+    -- either an input parse error or input has already been provided to this node
+    -- in either case, this is a failed match
+    Just _err -> return False
+  | otherwise -> return False
+
+
+-- | Match any node provided it's the final node in the query
 matchAnyAct :: NodeFinalAction
-matchAnyAct = NodeFinalAction $ \_ -> return True
+matchAnyAct = NodeFinalAction $ \remaining _node  -> return (null remaining)
 
 parseIdentQuery :: Parser NodeIdentQuery
 parseIdentQuery =
diff --git a/src/Pate/TraceTree.hs b/src/Pate/TraceTree.hs
index 182909aa..95b10317 100644
--- a/src/Pate/TraceTree.hs
+++ b/src/Pate/TraceTree.hs
@@ -110,6 +110,7 @@ module Pate.TraceTree (
   , NodeIdentQuery(..)
   , SomeTraceNode(..)
   , asChoice
+  , asInputChoice
   , forkTraceTreeHook
   , NodeFinalAction(..)
   , QueryResult(..)
@@ -118,7 +119,6 @@ module Pate.TraceTree (
   , giveChoiceInput
   , waitingForChoiceInput
   , chooseInput
-  , chooseInputLazy
   , chooseInputFromList
   ) where
 
@@ -317,7 +317,14 @@ instance Show NodeIdentQuery where
 
 data NodeFinalAction = NodeFinalAction
   {
-    finalAct :: forall l (k :: l). SomeTraceNode k -> IO Bool
+    -- | Takes the node at the terminal position as well as the remaining
+    -- elements in the query. For a typical menu choice, this requires that the
+    -- remaining query elements be empty (i.e. the given trace node is the result
+    -- of resolving the entire original query).
+    -- This allows a "choiceInput" node to be matched with a single query element
+    -- still remaining. This element is then parsed as input to the node, rather than
+    -- used to match a menu element.
+    finalAct :: forall l (k :: l). [NodeIdentQuery] -> SomeTraceNode k -> IO Bool
   }
 
 -- context plus final selection
@@ -347,6 +354,14 @@ asChoice (SomeTraceNode nm _ v) |
   = Just v
 asChoice _ = Nothing
 
+
+asInputChoice :: forall k. SomeTraceNode k -> Maybe (Some (InputChoice k))
+asInputChoice (SomeTraceNode nm _ v) |
+    Just Refl <- testEquality nm (knownSymbol @"choiceInput")
+  = Just v
+asInputChoice _ = Nothing
+
+
 data QueryResult k = QueryResult [NodeIdentQuery] (SomeTraceNode k)
 
 resolveQuery :: forall k.
@@ -360,23 +375,29 @@ resolveQuery (NodeQuery (q_outer:qs_outer) fin_outer) t_outer =
 
     go :: [NodeIdentQuery]-> NodeIdentQuery -> NodeQuery -> TraceTree k -> IO (Maybe (QueryResult k))
     go acc q (NodeQuery qs fin) (TraceTree t) = do
-      -- IO.putStrLn $ "go 1:" ++ show q
       result <- withIOList t $ \nodes -> do
-        -- IO.putStrLn "go 2"
         matches <- get_matches q nodes
         check_matches acc (NodeQuery qs fin) matches
       case result of
         Left{} -> return Nothing
         Right a -> return $ Just a
+    
+    -- Check all of the matches at the current query level against the finalization action.
+    check_final :: [NodeIdentQuery] -> NodeQuery -> [(NodeIdentQuery, SomeTraceNode k, TraceTree k)] -> IO (Maybe (QueryResult k))
+    check_final acc (NodeQuery remaining fin) ((matched_q, x,_):xs) = finalAct fin remaining x >>= \case
+      True -> return $ Just (QueryResult ((reverse (matched_q:acc)) ++ remaining) x)
+      False -> check_final acc (NodeQuery remaining fin) xs
+    check_final _ _ [] = return Nothing
 
     check_matches :: [NodeIdentQuery] -> NodeQuery -> [(NodeIdentQuery, SomeTraceNode k, TraceTree k)] -> IO (Maybe (QueryResult k))
-    check_matches acc (NodeQuery [] fin) ((q, x,_):xs) = finalAct fin x >>= \case
-      True -> return $ Just (QueryResult (reverse (q:acc)) x)
-      False -> check_matches acc (NodeQuery [] fin) xs
-    check_matches _ _ [] = return Nothing
-    check_matches acc (NodeQuery (q:qs) fin) ((matched_q, _,t):xs) = go (matched_q:acc) q (NodeQuery qs fin) t >>= \case
+    check_matches acc nodeQuery@(NodeQuery remaining fin) matches@((matched_q, _,t):xs) = check_final acc nodeQuery matches >>= \case
       Just result -> return $ Just result
-      Nothing -> check_matches acc (NodeQuery (q:qs) fin) xs
+      Nothing -> case remaining of
+        q:qs -> go (matched_q:acc) q (NodeQuery qs fin) t >>= \case
+          Just result -> return $ Just result
+          Nothing -> check_matches acc (NodeQuery (q:qs) fin) xs
+        [] -> return Nothing
+    check_matches _ _ [] = return Nothing
 
     get_matches :: NodeIdentQuery -> [Some (TraceTreeNode k)] -> IO [(NodeIdentQuery, SomeTraceNode k, TraceTree k)]
     get_matches q  nodes = do
@@ -1031,7 +1052,7 @@ data InputChoice (k :: l) nm where
   InputChoice :: (IsTraceNode k nm) =>
     { inputChoiceParse :: String -> Either InputChoiceError (TraceNodeLabel nm, TraceNodeType k nm)
     , inputChoicePut :: TraceNodeLabel nm -> TraceNodeType k nm -> IO Bool -- returns false if input has already been provided
-    , inputChoiceValue :: IO (Maybe (TraceNodeLabel nm, TraceNodeType k nm))
+    , inputChoiceValue :: IO (Maybe (TraceNodeLabel nm, TraceNodeType k nm)) 
     } -> InputChoice k nm
 
 waitingForChoiceInput :: InputChoice k nm -> IO Bool
@@ -1099,45 +1120,44 @@ chooseInput treenm parseInput = do
     Interactive nextChoiceIdent -> do
       newChoiceIdent <- liftIO $ nextChoiceIdent
       let status = NodeStatus StatusSuccess False (BlockedStatus (Set.singleton newChoiceIdent) Set.empty)
-      let statusFinal = NodeStatus StatusSuccess True (BlockedStatus Set.empty (Set.singleton newChoiceIdent))
+      let statusFinal = NodeStatus StatusSuccess False (BlockedStatus Set.empty (Set.singleton newChoiceIdent))
       c <- liftIO $ newEmptyMVar
-      let getValue = tryReadMVar c
-      let putValue lbl v = tryPutMVar c (lbl, v)
+      choice_lock <- liftIO $ newMVar False
+      let getValue = withMVar choice_lock $ \case
+            False -> return Nothing
+            True -> Just <$> readMVar c
+      let putValue lbl v = modifyMVar choice_lock $ \case
+            False -> do
+              putMVar c (lbl, v)
+              -- we need to mark this builder as unblocked before returning
+              -- to avoid having an intermediate state where the script runner sees
+              -- this node is still blocked on input that was just provided by the script
+              updateTreeStatus builder statusFinal
+              return (True, True)
+            True -> return (True, False)
       let ichoice = InputChoice @k @nm_choice parseInput putValue getValue
+      
       (lbl, v) <- withTracingLabel @"choiceInput" @k treenm (Some ichoice) $ do
         builder' <- getTreeBuilder
         liftIO $ updateTreeStatus builder' status
+        -- this blocks until 'putValue' sets the 'c' mvar, which either comes from
+        -- user input (i.e. via the Repl or GUI) or the script runner
         (lbl, v) <- liftIO $ readMVar c
         emitTraceLabel @nm_choice lbl v
+        -- ideally this inner builder is the one that would be unblocked by 'putValue', but
+        -- expressing this is a bit convoluted (i.e. because this builder doesn't exist when
+        -- defining 'putValue')
+        -- this just means that there is a temporary intermediate state where the 'choiceInput' node is
+        -- considered blocked waiting for input, while its parent node is unblocked
+        -- I don't think this is an issue in practice, although there are likely pathological
+        -- cases where a script will fail if it tries to match on the above element.
+        -- It's unclear why this would ever be useful, since this element is just a record of the value that
+        -- was parsed from the input.
         liftIO $ updateTreeStatus builder' statusFinal
         return (lbl, v)
       return $ Just (lbl, v)
     DefaultChoice -> return Nothing
 
-chooseInputLazy ::
-  forall nm_choice a b k m e.
-  IsTreeBuilder k e m =>
-  IsTraceNode k nm_choice =>
-  IO.MonadUnliftIO m =>
-  String ->
-  (String -> Either InputChoiceError (TraceNodeLabel nm_choice, TraceNodeType k nm_choice)) ->
-  (TraceNodeLabel nm_choice -> TraceNodeType k nm_choice -> b -> IO a) ->
-  m (LazyIOAction b a)
-chooseInputLazy treenm parseInput f = do
-  builder <- getTreeBuilder
-  case interactionMode builder of
-    Interactive{} -> do
-      c <- liftIO $ newEmptyMVar
-      let getValue = tryReadMVar c
-      let putValue lbl v = tryPutMVar c (lbl, v)
-      let isReady = not <$> isEmptyMVar c
-      let ichoice = InputChoice @k @nm_choice parseInput putValue getValue
-      emitTraceLabel @"choiceInput" @k treenm (Some ichoice)
-      return $ LazyIOAction (liftIO isReady) $ \inputVal -> getValue >>= \case
-        Just (lbl, v) -> Just <$> f lbl v inputVal
-        Nothing -> return Nothing
-    DefaultChoice{} -> return $ LazyIOAction (return False) (\_ -> return Nothing)
-
 choose ::
   forall nm_choice a k m e.
   IsTreeBuilder k e m =>
diff --git a/src/Pate/Verification/StrongestPosts.hs b/src/Pate/Verification/StrongestPosts.hs
index 62b45e03..d177b2c9 100644
--- a/src/Pate/Verification/StrongestPosts.hs
+++ b/src/Pate/Verification/StrongestPosts.hs
@@ -509,7 +509,7 @@ pickCutPoints pickMany msg inputs = go []
            , concBlk <- return $ PB.mkConcreteBlock blk PB.BlockEntryJump addr
            ]
       let opts = case pickMany && hasBin PBi.OriginalRepr picked && hasBin PBi.PatchedRepr picked of
-            True -> addr_opts ++ [("Finished Choosing", Nothing), ("", Nothing)]
+            True -> addr_opts ++ [("Finish Choosing", Nothing), ("", Nothing)]
             False -> addr_opts
       mres <- chooseInputFromList msg opts
       case mres of

From f2b6e3d317ea93d5a650e9bb2ab75e887131fa3a Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Wed, 24 Jul 2024 13:59:10 -0700
Subject: [PATCH 6/6] add '--prefer-text-input' flag, defaulting to original
 behavior

---
 src/Pate/CLI.hs                         | 5 +++++
 src/Pate/Config.hs                      | 4 ++++
 src/Pate/Verification/StrongestPosts.hs | 9 +++++++--
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/Pate/CLI.hs b/src/Pate/CLI.hs
index 5efdd9c5..0c5745b7 100644
--- a/src/Pate/CLI.hs
+++ b/src/Pate/CLI.hs
@@ -149,6 +149,7 @@ data CLIOptions = CLIOptions
   , noAssumeStackScope :: Bool
   , ignoreWarnings :: [String]
   , alwaysClassifyReturn :: Bool
+  , preferTextInput :: Bool
   } deriving (Eq, Ord, Show)
 
 printAtVerbosity
@@ -469,4 +470,8 @@ cliOptions = OA.info (OA.helper <*> parser)
     <*> OA.switch
          ( OA.long "always-classify-return"
          <> OA.help "Always resolve classifier failures by assuming function returns, if possible."
+         )
+    <*> OA.switch
+         ( OA.long "prefer-text-input"
+         <> OA.help "Prefer taking text input over multiple choice menus where possible."
          )
\ No newline at end of file
diff --git a/src/Pate/Config.hs b/src/Pate/Config.hs
index a0fd8ad8..9cfcf94d 100644
--- a/src/Pate/Config.hs
+++ b/src/Pate/Config.hs
@@ -287,6 +287,9 @@ data VerificationConfig validRepr =
     , cfgAlwaysClassifyReturn :: Bool
     -- ^ true if block classifier failures that can jump anywhere should be classified
     -- as returns without asking. 
+    , cfgPreferTextInput :: Bool
+    -- ^ modifies some menus to take string input instead of providing
+    -- a menu selection
     }
 
 
@@ -313,4 +316,5 @@ defaultVerificationCfg =
                      , cfgScriptPath = Nothing
                      , cfgIgnoreWarnings = []
                      , cfgAlwaysClassifyReturn = False
+                     , cfgPreferTextInput = False
                      }
diff --git a/src/Pate/Verification/StrongestPosts.hs b/src/Pate/Verification/StrongestPosts.hs
index d177b2c9..ddbf5e66 100644
--- a/src/Pate/Verification/StrongestPosts.hs
+++ b/src/Pate/Verification/StrongestPosts.hs
@@ -509,9 +509,14 @@ pickCutPoints pickMany msg inputs = go []
            , concBlk <- return $ PB.mkConcreteBlock blk PB.BlockEntryJump addr
            ]
       let opts = case pickMany && hasBin PBi.OriginalRepr picked && hasBin PBi.PatchedRepr picked of
-            True -> addr_opts ++ [("Finish Choosing", Nothing), ("", Nothing)]
+            True -> addr_opts ++ [("Finish Choosing", Nothing)]
             False -> addr_opts
-      mres <- chooseInputFromList msg opts
+      -- either generate a menu of alternatives or take text input, based on
+      -- the current configuration
+      mres <- asks (PCfg.cfgPreferTextInput . envConfig) >>= \case
+        True -> chooseInputFromList msg opts
+        False -> fmap Just $ 
+          choose @"()" msg $ \choice -> forM_ opts $ \(s,v) -> choice s () $ return v 
       case mres of
         Just (Just (Pair blk (PPa.WithBin bin addr))) -> do
           _ <- addIntraBlockCut addr blk