From 885e0a9eebe0b6f06d241e1aec32972343ad06ba Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Wed, 2 Oct 2024 16:32:23 -0700
Subject: [PATCH 1/6] Widening: promote WidenState to datatype that is threaded
 through widening loops

this is in preparation for collecting traces during widening, but
also cleans up some of the argument wrangling
---
 src/Pate/Verification/Widening.hs | 128 +++++++++++++++++-------------
 1 file changed, 73 insertions(+), 55 deletions(-)

diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index 5fcff5f6..d3c7de6c 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -1333,7 +1333,32 @@ tryWidenings (x:xs) =
     NoWideningRequired -> tryWidenings xs
     res -> return res
 
-type WidenState sym arch v = Either (AbstractDomain sym arch v) (WidenResult sym arch v)
+
+data WidenCase =
+    WidenCaseStart
+  | WidenCaseStep WidenKind
+  | WidenCaseErr String
+
+data WidenState sym arch v = WidenState
+  { 
+    stDomain :: AbstractDomain sym arch v
+  , stLocs :: WidenLocs sym arch
+  , stWidenCase :: WidenCase
+    -- ^ starting equivalence domain or accumulated widening result
+  , stTraces :: Map.Map (PL.SomeLocation sym arch) (CE.TraceEvents sym arch)
+    -- ^ traces collected during widening, indexed by which locations were widened at each step
+  }
+
+initWidenState :: AbstractDomain sym arch v -> WidenState sym arch v
+initWidenState d = WidenState d (WidenLocs Set.empty) WidenCaseStart Map.empty
+
+
+-- Compute a final 'WidenResult' from an (intermediate) 'WidenState' (surjective)
+widenStateToResult :: WidenState sym arch v -> WidenResult sym arch v
+widenStateToResult st = case stWidenCase st of
+  WidenCaseStart -> NoWideningRequired
+  WidenCaseStep k -> Widen k (stLocs st) (stDomain st)
+  WidenCaseErr msg -> WideningError msg (stLocs st) (stDomain st)
 
 -- | This gives a fixed amount of gas for traversing the
 --   widening loop. Setting this value too low seems to
@@ -1363,11 +1388,12 @@ widenPostcondition ::
   AbstractDomain sym arch v {- ^ postdomain -} ->
   EquivM sym arch (WidenResult sym arch v)
 widenPostcondition scope bundle preD postD0 = do
-  r <- withTracing @"debug" "widenPostcondition" $ withSym $ \sym -> do
+  st <- withTracing @"debug" "widenPostcondition" $ withSym $ \sym -> do
     eqCtx <- equivalenceContext
     traceBundle bundle "Entering widening loop"
     subTree @"domain" "Widening Steps" $
-      widenLoop sym localWideningGas eqCtx postD0 Nothing
+      widenLoop sym localWideningGas eqCtx (initWidenState postD0)
+  let r = widenStateToResult st
   case r of
     -- since widening was required, we show why it was needed
     Widen WidenEquality _ _postD1 -> withSym $ \sym -> do
@@ -1393,17 +1419,13 @@ widenPostcondition scope bundle preD postD0 = do
      PL.Location sym arch nm k ->
      W4.Pred sym ->
      EquivM_ sym arch (WidenState sym arch v)
-   widenOnce widenK (Gas i) mwb prevState loc goal = case prevState of
-     Right NoWideningRequired -> return prevState
-     Right (WideningError{}) -> return prevState
+   widenOnce widenK (Gas i) mwb prevState loc goal = case stWidenCase prevState of
+     WidenCaseErr{} -> return prevState
      _ -> startTimer $ withSym $ \sym -> do
        eqCtx <- equivalenceContext
-       (prevLocs, postD) <- case prevState of
-             Left prevDom -> return (mempty, prevDom)
-             --FIXME: we're dropping the widening Kind now since we're
-             --potentially doing multiple widenings in one iteration
-             Right (Widen _ locs prevDom) -> return (locs, prevDom)
-             -- NOTE: spurious missing case on some ghc versions
+       let 
+        this_loc = WidenLocs (Set.singleton (PL.SomeLocation loc))
+        postD = stDomain prevState
        curAsms <- currentAsm
        let emit r =
              withValid @() $ emitEvent (PE.SolverEvent (PS.simPair bundle) PE.EquivalenceProof r curAsms goal)
@@ -1425,12 +1447,12 @@ widenPostcondition scope bundle preD postD0 = do
            -- under analysis as inequivalent in the resulting domain
 
            case widenK of
-             WidenValue | Just (Some wb) <- mwb -> Right <$> dropValueLoc wb loc postD
+             WidenValue | Just (Some wb) <- mwb -> result <$> dropValueLoc wb loc postD
              WidenEquality ->
                case loc of
-                 PL.Cell c -> Right <$> widenCells [Some c] postD
-                 PL.Register r -> Right <$> widenRegs [Some r] postD
-                 PL.Unit -> return $ Right $ WideningError msg prevLocs postD
+                 PL.Cell c -> result <$> widenCells [Some c] postD
+                 PL.Register r -> result <$> widenRegs [Some r] postD
+                 PL.Unit -> return $ result $ WideningError msg this_loc postD
                  _ -> throwHere $ PEE.UnsupportedLocation
              _ -> panic Verifier "widenPostcondition" [ "Unexpected widening case"]
          Sat evalFn -> do
@@ -1443,17 +1465,17 @@ widenPostcondition scope bundle preD postD0 = do
                 let msg = unlines [ "Ran out of gas performing local widenings"
                                   , show (pretty ineqRes)
                                   ]
-                return $ Right $ WideningError msg prevLocs postD
+                return $ result $ WideningError msg this_loc postD
            else do
              -- The current execution does not satisfy the postcondition, and we have
              -- a counterexample.
              -- FIXME: postCondAsm doesn't exist anymore, but needs to be factored
              -- out still
-             res <- widenUsingCounterexample sym scope evalFn bundle eqCtx (W4.truePred sym) (PAD.absDomEq postD) preD prevLocs postD
+             res <- widenUsingCounterexample sym scope evalFn bundle eqCtx (W4.truePred sym) (PAD.absDomEq postD) preD postD
              case res of
                -- this location was made equivalent by a previous widening in this same loop
-               NoWideningRequired -> case prevState of
-                 Left{} ->  do
+               NoWideningRequired -> case stWidenCase prevState of
+                 WidenCaseStart ->  do
                    -- if we haven't performed any widenings yet, then this is an error
                    slice <- PP.simBundleToSlice scope bundle
                    ineqRes <- PP.getInequivalenceResult PEE.InvalidPostState
@@ -1463,9 +1485,15 @@ widenPostcondition scope bundle preD postD0 = do
                                      , show (pretty ineqRes)
                                      ]
                    
-                   return $ Right $ WideningError msg prevLocs postD
-                 Right{} -> return prevState
-               _ -> return $ Right res
+                   return $ result $ WideningError msg this_loc postD
+                 _ -> return prevState
+               _ -> return $ result res
+    where
+      result :: WidenResult sym arch v -> WidenState sym arch v
+      result r = case r of
+        NoWideningRequired -> prevState
+        WideningError err locs d -> prevState { stWidenCase = WidenCaseErr err, stLocs = locs <> stLocs prevState, stDomain = d }
+        Widen widenk locs d -> prevState { stWidenCase = WidenCaseStep widenk, stLocs = locs <> stLocs prevState, stDomain = d }
    
    -- The main widening loop. For now, we constrain it's iteration with a Gas parameter.
    -- In principle, I think this shouldn't be necessary, so we should revisit at some point.
@@ -1479,13 +1507,12 @@ widenPostcondition scope bundle preD postD0 = do
      sym ->
      Gas ->
      EquivContext sym arch ->
-     AbstractDomain sym arch v ->
-     Maybe (WidenResult sym arch v)
-     {- ^ A summary of any widenings that were done in previous iterations.
-          If @Nothing@, than no previous widenings have been performed. -} ->
-     NodeBuilderT '(sym,arch) "domain" (EquivM_ sym arch) (WidenResult sym arch v)
-   widenLoop sym (Gas i) eqCtx postD mPrevRes = subTraceLabel' PAD.Postdomain  (Some postD) $ \unlift ->
-     do
+     WidenState sym arch v
+     {- ^ A summary of any widenings that were done in previous iterations. -} ->
+     NodeBuilderT '(sym,arch) "domain" (EquivM_ sym arch) (WidenState sym arch v)
+   widenLoop sym (Gas i) eqCtx prevRes = 
+    let postD = stDomain prevRes
+    in subTraceLabel' PAD.Postdomain  (Some (stDomain prevRes)) $ \unlift -> do
         let (stO, stP) = PS.asStatePair scope (simOut bundle) PS.simOutState
         
         postVals <- PPa.forBinsC $ \bin -> do
@@ -1493,11 +1520,15 @@ widenPostcondition scope bundle preD postD0 = do
           st <- PPa.get bin $ PPa.PatchPair stO stP
           liftIO $ PAD.absDomainValsToPostCond sym eqCtx st Nothing vals
 
+        -- we reset the widen case so we can capture if this
+        -- step did anything
+        let res0 = prevRes { stWidenCase = WidenCaseStart }
+
         res2 <- case postVals of
           PPa.PatchPairSingle bin (Const valPost) ->
-            PL.foldLocation @sym @arch sym valPost (Left postD) (widenOnce WidenValue (Gas i) (Just (Some bin)))
+            PL.foldLocation @sym @arch sym valPost res0 (widenOnce WidenValue (Gas i) (Just (Some bin)))
           PPa.PatchPairC valPostO valPostP -> do
-            res1 <- PL.foldLocation @sym @arch sym valPostO (Left postD) (widenOnce WidenValue (Gas i) (Just (Some PBi.OriginalRepr)))
+            res1 <- PL.foldLocation @sym @arch sym valPostO res0 (widenOnce WidenValue (Gas i) (Just (Some PBi.OriginalRepr)))
             PL.foldLocation @sym @arch sym valPostP res1 (widenOnce WidenValue (Gas i) (Just (Some PBi.PatchedRepr)))
         
         -- for single-sided verification the equality condition is that the updated value is equal to the
@@ -1513,42 +1544,30 @@ widenPostcondition scope bundle preD postD0 = do
         -- was done in previous iterations (i.e., this is the first iteration)
         -- return `NoWideningRequired`.  Otherwise return the new abstract domain
         -- and a summary of the widenings we did.
-        case res of
-
+        case stWidenCase res of
           -- Some kind of error occured while widening.
-          Right er@(WideningError msg locs _postD') ->
+          WidenCaseErr msg ->
             do traceBundle bundle "== Widening error! =="
                traceBundle bundle msg
                traceBundle bundle "Partial widening at locations:"
-               traceBundle bundle (show locs)
+               traceBundle bundle (show (stLocs res))
 {-
                traceBundle bundle "===== PREDOMAIN ====="
                traceBundle bundle (show (PEE.ppEquivalenceDomain W4.printSymExpr (PS.specBody preD)))
                traceBundle bundle "===== POSTDOMAIN ====="
                traceBundle bundle (show (PEE.ppEquivalenceDomain W4.printSymExpr (PS.specBody postD')))
 -}
-               return er
+               return res
 
           -- In this iteration, no additional widening was done, and we can exit the loop.
           -- The ultimate result we return depends on if we did any widening steps in
-          -- previous iterations.
-          Right NoWideningRequired ->
-            case mPrevRes of
-              Nothing   -> return NoWideningRequired
-              Just prevRes -> return prevRes
-          -- no widening happened
-          Left{} ->
-            case mPrevRes of
-              Nothing   -> return NoWideningRequired
-              Just prevRes -> return prevRes
+          -- previous iterations (i.e. we restore the previous widen case)
+          WidenCaseStart -> return $ res { stWidenCase = stWidenCase prevRes }
           -- We had to do some widening in this iteration, so reenter the loop.
-          Right (Widen widenK locs postD') ->
+          WidenCaseStep{} ->
             do traceBundle bundle "== Found a widening, returning into the loop =="
-               traceBundle bundle (show locs)
-               let newlocs = case mPrevRes of
-                     Just (Widen _ prevLocs _) -> locs <> prevLocs
-                     _ -> locs
-               unlift $ widenLoop sym (Gas (i-1)) eqCtx postD' (Just $ Widen widenK newlocs postD')
+               traceBundle bundle (show (stLocs res))
+               unlift $ widenLoop sym (Gas (i-1)) eqCtx res
 
 
 -- | Refine a given 'AbstractDomainBody' to contain concrete values for the
@@ -1585,10 +1604,9 @@ widenUsingCounterexample ::
   W4.Pred sym ->
   PEE.EquivalenceDomain sym arch ->
   AbstractDomain sym arch v ->
-  WidenLocs sym arch {- ^ previous widening -}   ->
   AbstractDomain sym arch v ->
   EquivM sym arch (WidenResult sym arch v)
-widenUsingCounterexample sym scope evalFn bundle eqCtx postCondAsm postCondStatePred preD _prevLocs postD =
+widenUsingCounterexample sym scope evalFn bundle eqCtx postCondAsm postCondStatePred preD postD =
   tryWidenings
     [ -- First check for any disagreement in the constant values
       widenValues sym evalFn bundle postD

From 194e85dd6a825a07bde73cd2b77e416500fb85b0 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Fri, 4 Oct 2024 15:55:10 -0700
Subject: [PATCH 2/6] collect traces during widening and serialize with domain
 in common expression-binding environment

---
 pate.cabal                              |   1 +
 src/Pate/Monad.hs                       |  64 ++++++++++++-
 src/Pate/TraceCollection.hs             | 119 ++++++++++++++++++++++++
 src/Pate/Verification/AbstractDomain.hs |   8 ++
 src/Pate/Verification/Widening.hs       | 102 +++++++++++---------
 src/What4/JSON.hs                       |  46 +++++++--
 6 files changed, 286 insertions(+), 54 deletions(-)
 create mode 100644 src/Pate/TraceCollection.hs

diff --git a/pate.cabal b/pate.cabal
index 6d0a9793..e3a881c4 100644
--- a/pate.cabal
+++ b/pate.cabal
@@ -161,6 +161,7 @@ library
                        Pate.Script,
                        Pate.Timeout,
                        Pate.TraceTree,
+                       Pate.TraceCollection,
                        Pate.TraceConstraint,
                        Pate.ExprMappable,
                        What4.ExprHelpers,
diff --git a/src/Pate/Monad.hs b/src/Pate/Monad.hs
index a4fbcd1e..4050b8a4 100644
--- a/src/Pate/Monad.hs
+++ b/src/Pate/Monad.hs
@@ -103,6 +103,7 @@ module Pate.Monad
   , fnTrace
   , getWrappedSolver
   , joinPatchPred
+  , withSharedEnvEmit
   , module PME
   , atPriority, currentPriority, thisPriority)
   where
@@ -111,7 +112,7 @@ import           GHC.Stack ( HasCallStack, callStack )
 
 import           Control.Lens ( (&), (.~) )
 import qualified Control.Monad.Fail as MF
-import           Control.Monad (void)
+import           Control.Monad (void, forM)
 import           Control.Monad.IO.Class (liftIO)
 import qualified Control.Monad.IO.Unlift as IO
 import qualified Control.Concurrent as IO
@@ -120,6 +121,7 @@ import           Control.Monad.Catch
 import qualified Control.Monad.Reader as CMR
 import           Control.Monad.Reader ( asks, ask )
 import           Control.Monad.Except
+import           Control.Monad.State (StateT(..), get, put, evalStateT, modify, execStateT)
 
 import           Data.Maybe ( fromMaybe )
 import qualified Data.Map as M
@@ -197,6 +199,9 @@ import Data.Functor.Const (Const(..))
 import Unsafe.Coerce (unsafeCoerce)
 import Debug.Trace
 import Data.List
+import Data.Parameterized.SymbolRepr
+import qualified What4.JSON as W4S
+import qualified Data.Aeson as JSON
 
 atPriority :: 
   NodePriority ->
@@ -402,7 +407,62 @@ printExprTruncated (ExprLabel lbl) (Some e) =
     [] -> pfx
     [a] -> pfx <> PP.pretty a
     (a:as) -> pfx <> PP.pretty a <> ".." <> PP.pretty (last as)
-              
+
+type IsSerializableNode sym arch nm = (IsTraceNode '(sym, arch) nm, W4S.W4Serializable sym (TraceNodeLabel nm), W4S.W4Serializable sym (TraceNodeType '(sym, arch) nm))
+
+data W4SerializableNode sym arch nm = IsSerializableNode sym arch nm => 
+  W4SerializableNode (SymbolRepr nm) (TraceNodeLabel nm) (TraceNodeType '(sym, arch) nm)
+
+-- | List of trace nodes (of any type) that share the same expression binding
+--   environment when serialized
+newtype SharedExprEnv sym arch = SharedExprEnv [Some (W4SerializableNode sym arch)]
+
+instance ValidSymArch sym arch => IsTraceNode '(sym,arch) "shared_env" where
+  type TraceNodeType '(sym,arch) "shared_env" = SharedExprEnv sym arch
+  type TraceNodeLabel "shared_env" = String
+
+  prettyNode top_lbl (SharedExprEnv nds) = 
+    PP.vsep $ (PP.pretty top_lbl <> ": "):(map (\(Some (W4SerializableNode (_ :: SymbolRepr nm) lbl nd)) -> prettyNode @_ @'(sym,arch) @nm lbl nd) nds)
+
+  jsonNode sym top_lbl (SharedExprEnv nds) = do
+    (\f -> evalStateT f W4S.emptyExprCache) $ do
+      contents <- forM nds $ \(Some (W4SerializableNode nm lbl nd)) -> do
+        cache0 <- get
+        (lbl_json, cache1) <- liftIO $ W4S.w4ToJSONWithCache sym cache0 lbl
+        (nd_json, cache2) <- liftIO $ W4S.w4ToJSONWithCache sym cache1 nd
+        put cache2
+        return $ JSON.object [ "type" JSON..= symbolRepr nm, "label" JSON..= lbl_json, "value" JSON..= nd_json ]
+      cache <- get
+      env <- liftIO $ W4S.cacheToEnv sym cache
+      env_json <- liftIO $ W4S.serializeExprEnv sym env
+      return $ JSON.object [ "name" JSON..= top_lbl, "shared_env" JSON..= env_json, "contents" JSON..= contents]
+
+-- | Emit multiple values to the tracetree with a common expression binding environment
+--   This will ultimately result in a single composite value, where all of the given values are serialized
+--   using the same expression binding cache.
+--
+--
+--   FIXME: this is a bit of a gross hack to work around the fact that we can't supply
+--   an expression binding environment when emitting a value via 'emitTrace'.
+--   What we need is a common datatype for expression binding environments that both
+--   TraceTree and W4Serializable use so we can do this more sensibly (i.e. have a IsTreeBuilder primitive that
+--   collects all of the inner traces and builds a common binding environment from them)
+--
+--   For now we can at least abstract away the details of this packaging by just taking
+--   an arbitrary monadic operation that collects the individual traces.
+
+withSharedEnvEmit ::
+  forall sym arch.
+  String ->
+  (forall m. Monad m => (forall nm. (IsSerializableNode sym arch nm) => SymbolRepr nm -> TraceNodeLabel nm -> TraceNodeType '(sym,arch) nm -> m ()) -> m ()) ->
+  EquivM sym arch ()
+withSharedEnvEmit top_nm f = do
+  env <- execStateT (f g) (SharedExprEnv [])
+  emitTraceLabel @"shared_env" @'(sym,arch) top_nm env
+  where 
+    g :: forall (nm :: Symbol). IsSerializableNode sym arch nm => SymbolRepr nm -> TraceNodeLabel nm -> TraceNodeType '(sym,arch) nm -> StateT (SharedExprEnv sym arch) (EquivM_ sym arch) ()
+    g nm lbl v = modify $ \(SharedExprEnv vs) -> SharedExprEnv ((Some (W4SerializableNode nm lbl v)):vs)
+
 withBinary ::
   forall bin sym arch a.
   PBi.KnownBinary bin =>
diff --git a/src/Pate/TraceCollection.hs b/src/Pate/TraceCollection.hs
new file mode 100644
index 00000000..affe171d
--- /dev/null
+++ b/src/Pate/TraceCollection.hs
@@ -0,0 +1,119 @@
+
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+
+module Pate.TraceCollection
+  ( TraceCollection
+  , empty
+  , lookupReg
+  , insertReg
+  , lookupCell
+  , insertCell
+  ) where
+
+import           Prelude hiding ( lookup )
+
+import qualified Prettyprinter as PP
+
+import qualified Data.Set as Set
+import           Data.Set ( Set )
+
+import qualified Data.Map as Map
+import           Data.Map ( Map )
+
+import qualified Data.Text as T
+
+import qualified Data.Macaw.CFG as MM
+
+import           Data.Parameterized.Classes
+import           Data.Parameterized.Some
+
+import qualified Pate.Arch as PA
+import           Pate.TraceTree
+import qualified Pate.MemCell as PMC
+import qualified Pate.Solver as PSo
+import qualified Pate.Verification.StrongestPosts.CounterExample as CE
+
+import qualified What4.JSON as W4S
+import qualified Data.Aeson as JSON
+
+data TraceCollection sym arch = TraceCollection
+  { 
+    trAllTraces :: [CE.TraceEvents sym arch]
+  , trTraceMapRegs :: Map (Some (MM.ArchReg arch)) (Set Int)
+  -- ^ mapping from registers into indexes in to the list of all traces
+  , trTraceMapCells :: Map (Some (PMC.MemCell sym arch)) (Set Int)
+  -- ^ mapping from memory cells into indexes in to the list of all traces
+  }
+
+empty :: TraceCollection sym arch
+empty = TraceCollection [] Map.empty Map.empty
+
+insertReg ::
+  PA.ValidArch arch =>
+  MM.ArchReg arch tp ->
+  CE.TraceEvents sym arch ->
+  TraceCollection sym arch ->
+  TraceCollection sym arch
+insertReg reg tr trcol = trcol 
+  { trAllTraces = tr:(trAllTraces trcol)
+  , trTraceMapRegs = Map.insertWith Set.union (Some reg) (Set.singleton (length (trAllTraces trcol))) (trTraceMapRegs trcol)
+  }
+
+insertCell ::
+  PSo.ValidSym sym =>
+  PMC.MemCell sym arch w ->
+  CE.TraceEvents sym arch ->
+  TraceCollection sym arch ->
+  TraceCollection sym arch
+insertCell cell tr trcol = trcol 
+  { trAllTraces = tr:(trAllTraces trcol)
+  , trTraceMapCells = Map.insertWith Set.union (Some cell) (Set.singleton (length (trAllTraces trcol))) (trTraceMapCells trcol)
+  }
+
+lookupReg ::
+  PA.ValidArch arch =>
+  TraceCollection sym arch ->
+  MM.ArchReg arch tp ->
+  [CE.TraceEvents sym arch]
+lookupReg trcol reg = case Map.lookup (Some reg) (trTraceMapRegs trcol) of
+  Just idxs -> map (\i -> (trAllTraces trcol) !! i) (Set.toList idxs)
+  Nothing -> []
+
+lookupCell ::
+  (PSo.ValidSym sym, PA.ValidArch arch) =>
+  TraceCollection sym arch ->
+  PMC.MemCell sym arch w ->
+  [CE.TraceEvents sym arch]
+lookupCell trcol cell = case Map.lookup (Some cell) (trTraceMapCells trcol) of
+  Just idxs -> map (\i -> (trAllTraces trcol) !! i) (Set.toList idxs)
+  Nothing -> []
+
+instance (PA.ValidArch arch, PSo.ValidSym sym) => W4S.W4Serializable sym (TraceCollection sym arch) where
+  w4Serialize (TraceCollection allTraces regs cells) = do
+    W4S.object 
+      [ "all_traces" W4S..= allTraces
+      , "trace_map_regs" W4S..= regs
+      , "trace_map_cells" W4S..= cells
+      ]
+
+instance forall sym arch. (PA.ValidArch arch, PSo.ValidSym sym) => IsTraceNode '(sym,arch) "trace_collection" where
+  type TraceNodeType '(sym,arch) "trace_collection" = TraceCollection sym arch
+  type TraceNodeLabel "trace_collection" = T.Text
+
+  prettyNode nm _tc = "Trace Collection: " <> PP.pretty nm
+  nodeTags = mkTags @'(sym,arch) @"trace_collection" [Summary, Simplified]
+  jsonNode sym lbl v = do
+    v_json <- W4S.w4ToJSON sym v
+    return $ JSON.object [ "name"  JSON..= lbl , "trace_collection" JSON..= v_json  ]
\ No newline at end of file
diff --git a/src/Pate/Verification/AbstractDomain.hs b/src/Pate/Verification/AbstractDomain.hs
index acea8105..fad5f896 100644
--- a/src/Pate/Verification/AbstractDomain.hs
+++ b/src/Pate/Verification/AbstractDomain.hs
@@ -141,6 +141,9 @@ instance forall sym arch v. (PSo.ValidSym sym, PA.ValidArch arch) => W4S.W4Seria
     eq_dom <- W4S.w4Serialize (absDomEq abs_dom)
     return $ JSON.object [ "eq_domain" .= eq_dom, "val_domain" .= JSON.String "TODO" ]
 
+instance forall sym arch. (PSo.ValidSym sym, PA.ValidArch arch) => W4S.W4SerializableF sym (AbstractDomain sym arch)
+
+
 -- | Restrict an abstract domain to a single binary.
 singletonDomain ::
   PPa.PatchPairM m =>
@@ -888,6 +891,11 @@ ppDomainKind = \case
   Postdomain -> "Intermediate postdomain"
   ExternalPostDomain -> "Postdomain"
 
+instance JSON.ToJSON DomainKind where
+  toJSON dk = JSON.toJSON (show dk)
+
+instance W4S.W4Serializable sym DomainKind
+
 instance (PA.ValidArch arch, PSo.ValidSym sym) => IsTraceNode '(sym,arch) "domain" where
   type TraceNodeType '(sym,arch) "domain" = Some (AbstractDomain sym arch)
   type TraceNodeLabel "domain" = DomainKind
diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index d3c7de6c..51ee550b 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -100,6 +100,7 @@ import qualified Pate.PatchPair as PPa
 import qualified Pate.SimState as PS
 import qualified Pate.SimulatorRegisters as PSR
 import qualified Pate.Config as PC
+import qualified Pate.TraceCollection as PTc
 
 import           Pate.Verification.PairGraph
 import qualified Pate.Verification.ConditionalEquiv as PVC
@@ -481,30 +482,27 @@ addRefinementChoice nd gr0 = withTracing @"message" ("Modify Proof Node: " ++ sh
         return $ emptyWorkList gr4
 
 tryWithAsms :: 
-  [(W4.Pred sym, PEE.InnerEquivalenceError arch)] -> 
-  (SymGroundEvalFn sym -> EquivM_ sym arch a) ->
+  [(W4.Pred sym, PEE.InnerEquivalenceError arch)] ->
+  W4.Pred sym ->
+  (SatResult (SymGroundEvalFn sym) () -> EquivM_ sym arch a) ->
   EquivM sym arch a
-tryWithAsms ((asm, err):asms) f = do
-  mresult <- withSatAssumption (PAs.fromPred asm) $ tryWithAsms asms f
+tryWithAsms ((asm, err):asms) p f = do
+  mresult <- withSatAssumption (PAs.fromPred asm) $ tryWithAsms asms p f
   case mresult of
     Just a -> return a
-    Nothing -> emitWarning err >> tryWithAsms asms f
-tryWithAsms [] f = withSym $ \sym -> goalSat "tryWithAsms" (W4.truePred sym) $ \res -> case res of
-  Unsat _ -> throwHere PEE.InvalidSMTModel
-  Unknown -> throwHere PEE.InconclusiveSAT
-  Sat evalFn -> f evalFn
+    Nothing -> emitWarning err >> tryWithAsms asms p f
+tryWithAsms [] p f = goalSat "tryWithAsms" p f
 
-getSomeGroundTrace ::
-  PS.SimScope sym arch v ->
+-- | Check the predicate for satisfiability, attempting to
+--   constrain the model to valid pointers and zero-offset stacks
+--   if possible.
+withTraceAssumptions ::
   SimBundle sym arch v ->
-  AbstractDomain sym arch v ->
-  Maybe (StatePostCondition sym arch v) ->
-  EquivM sym arch (CE.TraceEvents sym arch)
-getSomeGroundTrace scope bundle preD postCond = withSym $ \sym -> do
+  W4.Pred sym ->
+  (SatResult (SymGroundEvalFn sym) () -> EquivM_ sym arch a) ->
+  EquivM sym arch a
+withTraceAssumptions bundle p f = withSym $ \sym -> do
   (_, ptrAsserts) <- PVV.collectPointerAssertions bundle
-  
-  -- try to ground the model with a zero stack base, so calculated stack offsets
-  -- are the same as stack slots
   stacks_zero <- PPa.catBins $ \bin -> do
     in_ <- PPa.get bin (PS.simIn bundle)
     let stackbase = PS.unSE $ PS.simStackBase (PS.simInState in_) 
@@ -513,12 +511,21 @@ getSomeGroundTrace scope bundle preD postCond = withSym $ \sym -> do
 
   ptrAsserts_pred <- PAs.toPred sym (stacks_zero <> ptrAsserts)
 
-  tr <- tryWithAsms 
+  tryWithAsms
     [ (ptrAsserts_pred, PEE.RequiresInvalidPointerOps)
-    ] $ \evalFn -> 
-      getTraceFromModel scope evalFn bundle preD postCond
+    ] p f
 
-  return tr
+getSomeGroundTrace ::
+  PS.SimScope sym arch v ->
+  SimBundle sym arch v ->
+  AbstractDomain sym arch v ->
+  Maybe (StatePostCondition sym arch v) ->
+  EquivM sym arch (CE.TraceEvents sym arch)
+getSomeGroundTrace scope bundle preD postCond = withSym $ \sym -> do
+  withTraceAssumptions bundle (W4.truePred sym) $ \case
+    Unsat{} -> throwHere PEE.InvalidSMTModel
+    Unknown -> throwHere PEE.InconclusiveSAT
+    Sat evalFn -> getTraceFromModel scope evalFn bundle preD postCond
 
 getTraceFootprint ::
   forall sym arch v.
@@ -1339,18 +1346,21 @@ data WidenCase =
   | WidenCaseStep WidenKind
   | WidenCaseErr String
 
+
 data WidenState sym arch v = WidenState
   { 
     stDomain :: AbstractDomain sym arch v
   , stLocs :: WidenLocs sym arch
   , stWidenCase :: WidenCase
     -- ^ starting equivalence domain or accumulated widening result
-  , stTraces :: Map.Map (PL.SomeLocation sym arch) (CE.TraceEvents sym arch)
-    -- ^ traces collected during widening, indexed by which locations were widened at each step
+  , stTracesEq :: PTc.TraceCollection sym arch
+    -- ^ collected traces for equality widening steps
+  , stTracesVal :: PTc.TraceCollection sym arch
+    --- ^ collected traces for value widening steps
   }
 
 initWidenState :: AbstractDomain sym arch v -> WidenState sym arch v
-initWidenState d = WidenState d (WidenLocs Set.empty) WidenCaseStart Map.empty
+initWidenState d = WidenState d (WidenLocs Set.empty) WidenCaseStart PTc.empty PTc.empty
 
 
 -- Compute a final 'WidenResult' from an (intermediate) 'WidenState' (surjective)
@@ -1393,23 +1403,14 @@ widenPostcondition scope bundle preD postD0 = do
     traceBundle bundle "Entering widening loop"
     subTree @"domain" "Widening Steps" $
       widenLoop sym localWideningGas eqCtx (initWidenState postD0)
-  let r = widenStateToResult st
-  case r of
-    -- since widening was required, we show why it was needed
-    Widen WidenEquality _ _postD1 -> withSym $ \sym -> do
-      eqCtx <- equivalenceContext
-      eqPost <- liftIO $ PEq.getPostdomain sym scope bundle eqCtx (PAD.absDomEq preD) (PAD.absDomEq postD0)
-      eqPost_pred <- liftIO $ postCondPredicate sym eqPost
-      withTracing @"message" "Equivalence Counter-example" $ do
-        not_eqPost_pred <- liftIO $ W4.notPred sym eqPost_pred
-        mres <- withSatAssumption (PAs.fromPred not_eqPost_pred) $ do
-          res <- getSomeGroundTrace scope bundle preD (Just eqPost)
-          emitTrace @"trace_events" res
-        case mres of
-          Just () -> return ()
-          Nothing -> emitWarning (PEE.WideningError "Couldn't find widening counter-example")
-        return r
-    _ -> return r
+    
+  case stWidenCase st of
+    WidenCaseStep _ -> withSharedEnvEmit "Equivalence Counter-example Traces" $ \emit -> do
+      emit (CT.knownSymbol @"trace_collection") "Equality" (stTracesEq st)
+      emit (CT.knownSymbol @"trace_collection") "Value" (stTracesVal st)
+      emit (CT.knownSymbol @"domain") PAD.Postdomain (Some (stDomain st))
+    _ -> return ()
+  return $ widenStateToResult st
  where
    widenOnce ::
      WidenKind ->
@@ -1435,7 +1436,7 @@ widenPostcondition scope bundle preD postD0 = do
        
        --(not_goal', ptrAsms) <- PVV.collectPointerAssertions not_goal
        emitTraceLabel @"expr" "goal" (Some goal)
-       goalSat "prove postcondition" not_goal $ \case
+       withTraceAssumptions bundle not_goal $ \case
          Unsat _ -> do
            emit PE.SolverSuccess
            return prevState
@@ -1487,8 +1488,23 @@ widenPostcondition scope bundle preD postD0 = do
                    
                    return $ result $ WideningError msg this_loc postD
                  _ -> return prevState
+               Widen widenk (WidenLocs locs) d -> do
+                -- FIXME: should we make a post condition here?
+                tr <- getTraceFromModel scope evalFn bundle preD Nothing
+                let st' = foldr (addLoc widenk tr) prevState locs 
+                return $ st' { stWidenCase = WidenCaseStep widenk, stLocs = WidenLocs locs <> stLocs prevState, stDomain = d }
                _ -> return $ result res
     where
+      addLoc :: WidenKind -> CE.TraceEvents sym arch -> PL.SomeLocation sym arch ->  WidenState sym arch v -> WidenState sym arch v
+      addLoc wk te (PL.SomeLocation l)  st = case (wk,l) of
+        (WidenEquality, PL.Register r) -> st { stTracesEq = PTc.insertReg r te (stTracesEq st) }
+        (WidenEquality, PL.Cell c) -> st { stTracesEq = PTc.insertCell c te (stTracesEq st) }
+        (WidenValue, PL.Register r) -> st { stTracesVal = PTc.insertReg r te (stTracesVal st) }
+        (WidenValue, PL.Cell c) -> st { stTracesVal = PTc.insertCell c te (stTracesVal st) }
+        -- other kinds of locations we can ignore for now: since this is just for reporting purposes we
+        -- only need to index the traces by locations the user actually knows about
+        _ -> st
+
       result :: WidenResult sym arch v -> WidenState sym arch v
       result r = case r of
         NoWideningRequired -> prevState
diff --git a/src/What4/JSON.hs b/src/What4/JSON.hs
index f942c84b..a23a9344 100644
--- a/src/What4/JSON.hs
+++ b/src/What4/JSON.hs
@@ -32,12 +32,17 @@ module What4.JSON
   , (.=~)
   , ExprEnv
   , mergeEnvs
+  , ExprCache
+  , cacheToEnv
+  , emptyExprCache
+  , w4ToJSONWithCache
   , W4Deserializable(..)
   , jsonToW4
   , readJSON
   , (.:)
   , SymDeserializable(..)
   , symDeserializable
+  , serializeExprEnv
   ) where
 
 import           Control.Monad.State (MonadState (..), State, modify, evalState, runState)
@@ -46,6 +51,8 @@ import qualified Data.Map.Ordered as OMap
 import           Data.List ( stripPrefix )
 import           Data.Maybe ( mapMaybe, catMaybes )
 import           Data.Map (Map)
+import           Data.Set (Set)
+import qualified Data.Set as Set
 import qualified Data.Map.Merge.Strict as Map
 import           Data.Text (Text)
 import           Data.Data (Proxy(..), Typeable)
@@ -115,6 +122,9 @@ instance MonadState (ExprCache sym) (W4S sym) where
 class W4Serializable sym a where
   w4Serialize :: a -> W4S sym JSON.Value
 
+  default w4Serialize :: JSON.ToJSON a => a -> W4S sym JSON.Value
+  w4Serialize a = return $ JSON.toJSON a
+
 w4SerializeString :: Show a => a -> W4S sym JSON.Value
 w4SerializeString s = return $ JSON.String (T.pack (show s))
 
@@ -127,6 +137,14 @@ w4ToJSON sym a = do
 
 newtype ExprEnv sym = ExprEnv (Map Integer (Some (W4.SymExpr sym)))
 
+-- | Serialize an expression binding environment
+serializeExprEnv :: forall sym. SerializableExprs sym => sym -> ExprEnv sym -> IO JSON.Value
+serializeExprEnv sym (ExprEnv env) = JSON.toJSON <$> do
+  forM (Map.toList env) $ \(i, Some (e :: W4.SymExpr sym tp)) -> do
+    e_json <- withSerializable (Proxy @sym) (Proxy @(W4.SymExpr sym)) (Proxy @tp) $ w4ToJSON sym e
+    return $ JSON.object [ "symbolic_ident" JSON..= i, "symbolic_expr" JSON..= e_json]
+
+
 mergeEnvs :: TestEquality (W4.SymExpr sym) => ExprEnv sym -> ExprEnv sym -> ExprEnv sym
 mergeEnvs (ExprEnv env1) (ExprEnv env2) = ExprEnv (
   Map.merge
@@ -157,6 +175,18 @@ w4ToJSONEnv sym a = do
   eenv <- cacheToEnv sym c
   return $ (v, eenv)
 
+emptyExprCache :: ExprCache sym
+emptyExprCache = ExprCache Map.empty
+
+w4ToJSONWithCache :: forall sym a. (W4.IsExprBuilder sym, SerializableExprs sym) => W4Serializable sym a => sym -> ExprCache sym -> a -> IO (JSON.Value, ExprCache sym)
+w4ToJSONWithCache sym init_cache a = do
+  cacheRef <- IO.newIORef init_cache
+  W4S f <- return $ w4Serialize @sym a
+  let env = W4SEnv cacheRef sym True
+  v <- runReaderT f env
+  c <- IO.readIORef cacheRef
+  return $ (v, c)
+
 class W4SerializableF sym (f :: k -> Type) where
   withSerializable :: Proxy sym -> p f -> q tp -> (W4Serializable sym (f tp) => a) -> a
 
@@ -243,8 +273,7 @@ instance (W4Serializable sym x) => W4Serializable sym [x] where
     xs_json <- mapM w4Serialize xs
     return $ JSON.toJSON xs_json
 
-instance W4Serializable sym JSON.Value where
-  w4Serialize = return
+instance W4Serializable sym JSON.Value
 
 instance (W4Serializable sym x, W4Serializable sym y) => W4Serializable sym (Map x y) where
   w4Serialize x = do
@@ -261,14 +290,13 @@ instance (SerializableExprs sym, W4Serializable sym x) => W4Serializable sym (W4
           W4P.PredDisjRepr -> "disj"
     object [ "predmap" .= objs, "kind" .= repr ]
 
-instance W4Serializable sym Integer where
-  w4Serialize i = return $ JSON.toJSON i
-
-instance W4Serializable sym Text where
-  w4Serialize i = return $ JSON.toJSON i
+instance W4Serializable sym Integer
+instance W4Serializable sym Int
+instance W4Serializable sym Text
+instance W4Serializable sym Bool
 
-instance W4Serializable sym Bool where
-  w4Serialize i = return $ JSON.toJSON i
+instance W4Serializable sym a => W4Serializable sym (Set a) where
+  w4Serialize a = w4Serialize (Set.toList a)
 
 instance W4Serializable sym a => W4Serializable sym (Maybe a) where
   w4Serialize = \case

From 78c8db65dffe819cdeca0ff74de521138cd3b42b Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Mon, 7 Oct 2024 10:38:22 -0700
Subject: [PATCH 3/6] TraceCollection: add 'insert' function to interface

this allows adding for associating multiple locations
with a single trace (that is not duplicated in the
underlying structure)
---
 src/Pate/TraceCollection.hs       | 91 ++++++++++++++++++++++++++++++-
 src/Pate/Verification/Widening.hs | 20 ++++---
 2 files changed, 101 insertions(+), 10 deletions(-)

diff --git a/src/Pate/TraceCollection.hs b/src/Pate/TraceCollection.hs
index affe171d..96511463 100644
--- a/src/Pate/TraceCollection.hs
+++ b/src/Pate/TraceCollection.hs
@@ -1,4 +1,15 @@
-
+{-|
+Module           : Pate.TraceCollection
+Copyright        : (c) Galois, Inc 2024
+Maintainer       : Daniel Matichuk <dmatichuk@galois.com>
+
+Specialized map that relates memory cells (see 'Pate.MemCell') and registers
+to traces. Used during widening (see 'Pate.Verification.Widening') to associate
+location that are widened in an equivalence domain to a trace that demonstrates
+why the widening was necessary (i.e. counter-example for how that location could
+be made inequivalent).
+
+-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -20,12 +31,15 @@ module Pate.TraceCollection
   , insertReg
   , lookupCell
   , insertCell
+  , insert
+  , lookup
   ) where
 
 import           Prelude hiding ( lookup )
 
 import qualified Prettyprinter as PP
 
+import           Data.Maybe
 import qualified Data.Set as Set
 import           Data.Set ( Set )
 
@@ -36,7 +50,7 @@ import qualified Data.Text as T
 
 import qualified Data.Macaw.CFG as MM
 
-import           Data.Parameterized.Classes
+import           Data.Parameterized.Classes()
 import           Data.Parameterized.Some
 
 import qualified Pate.Arch as PA
@@ -48,6 +62,11 @@ import qualified Pate.Verification.StrongestPosts.CounterExample as CE
 import qualified What4.JSON as W4S
 import qualified Data.Aeson as JSON
 
+
+-- | A map that associates locations ('MM.ArchReg' or 'PMC.MemCell') with traces
+--  ('CE.TraceEvents'). Each location is mapped to a set of indexes into
+--  a list of traces. These indexes are used during serialization ('W4S.W4Serializable') to
+--  avoid duplication when one trace is shared by multiple locations.
 data TraceCollection sym arch = TraceCollection
   { 
     trAllTraces :: [CE.TraceEvents sym arch]
@@ -60,6 +79,7 @@ data TraceCollection sym arch = TraceCollection
 empty :: TraceCollection sym arch
 empty = TraceCollection [] Map.empty Map.empty
 
+-- | Add a new trace to the set of traces associated with the given 'MM.ArchReg'
 insertReg ::
   PA.ValidArch arch =>
   MM.ArchReg arch tp ->
@@ -71,6 +91,7 @@ insertReg reg tr trcol = trcol
   , trTraceMapRegs = Map.insertWith Set.union (Some reg) (Set.singleton (length (trAllTraces trcol))) (trTraceMapRegs trcol)
   }
 
+-- | Add a new trace to the set of traces associated with the given 'PMC.MemCell'
 insertCell ::
   PSo.ValidSym sym =>
   PMC.MemCell sym arch w ->
@@ -82,6 +103,7 @@ insertCell cell tr trcol = trcol
   , trTraceMapCells = Map.insertWith Set.union (Some cell) (Set.singleton (length (trAllTraces trcol))) (trTraceMapCells trcol)
   }
 
+-- | Get all traces associated with the given 'MM.ArchReg'
 lookupReg ::
   PA.ValidArch arch =>
   TraceCollection sym arch ->
@@ -91,6 +113,7 @@ lookupReg trcol reg = case Map.lookup (Some reg) (trTraceMapRegs trcol) of
   Just idxs -> map (\i -> (trAllTraces trcol) !! i) (Set.toList idxs)
   Nothing -> []
 
+-- | Get all traces associated with the given 'PMC.MemCell'
 lookupCell ::
   (PSo.ValidSym sym, PA.ValidArch arch) =>
   TraceCollection sym arch ->
@@ -100,6 +123,70 @@ lookupCell trcol cell = case Map.lookup (Some cell) (trTraceMapCells trcol) of
   Just idxs -> map (\i -> (trAllTraces trcol) !! i) (Set.toList idxs)
   Nothing -> []
 
+-- | Add a single trace to the set of traces associated with the given
+--  list of registers and memory locations. Note that although this
+--  is functionally equivalent to folding via 'insertReg' and 'insertCell',
+--  the resulting JSON from serialization (via 'W4S.W4Serializable.w4Serialize')
+--  only contains one copy of the given trace.
+insert ::
+  PSo.ValidSym sym =>
+  PA.ValidArch arch =>
+  [Some (MM.ArchReg arch)] ->
+  [Some (PMC.MemCell sym arch)] ->
+  CE.TraceEvents sym arch ->
+  TraceCollection sym arch ->
+  TraceCollection sym arch
+insert regs cells tr trcol = trcol
+  { trAllTraces = tr:(trAllTraces trcol)
+  , trTraceMapRegs =
+      foldr (\reg -> Map.insertWith Set.union reg (Set.singleton idx)) (trTraceMapRegs trcol) regs
+  , trTraceMapCells =
+      foldr (\cell -> Map.insertWith Set.union cell (Set.singleton idx)) (trTraceMapCells trcol) cells
+  }
+  where
+    idx = length (trAllTraces trcol)
+
+-- | Find all traces associated with the given list of registers and memory locations
+--   (i.e. each trace is associated with at least one of the given locations).
+--   Traces that are associated with multiple locations (i.e. added with 'insert') only
+--   occur once in the result.
+lookup ::
+  PSo.ValidSym sym =>
+  PA.ValidArch arch =>
+  [Some (MM.ArchReg arch)] ->
+  [Some (PMC.MemCell sym arch)] ->
+  TraceCollection sym arch ->
+  [CE.TraceEvents sym arch]
+lookup regs cells trcol = let
+  reg_idxs = Set.unions $ map (\reg -> fromMaybe Set.empty $ Map.lookup reg (trTraceMapRegs trcol)) regs
+  cell_idxs = Set.unions $ map (\cell -> fromMaybe Set.empty $ Map.lookup cell (trTraceMapCells trcol)) cells
+  in map (\i -> (trAllTraces trcol) !! i) (Set.toList (Set.union reg_idxs cell_idxs))
+
+{-
+Not used a the moment, so left commented out to avoid cluttering the interface.
+
+toList ::
+  forall sym arch.
+  TraceCollection sym arch ->
+  [(([Some (MM.ArchReg arch)], [Some (PMC.MemCell sym arch)]), CE.TraceEvents sym arch)]
+toList trcol = map go [0..(length (trAllTraces trcol))]
+  where
+    go :: Int -> (([Some (MM.ArchReg arch)], [Some (PMC.MemCell sym arch)]), CE.TraceEvents sym arch)
+    go i = let
+      tr = trAllTraces trcol !! i
+      regs = Map.keys $ Map.filter (Set.member i) (trTraceMapRegs trcol)
+      cells = Map.keys $ Map.filter (Set.member i) (trTraceMapCells trcol)
+      in ((regs, cells), tr)
+
+fromList ::
+  forall sym arch.
+  PSo.ValidSym sym =>
+  PA.ValidArch arch =>
+  [(([Some (MM.ArchReg arch)], [Some (PMC.MemCell sym arch)]), CE.TraceEvents sym arch)] ->
+  TraceCollection sym arch
+fromList trs = foldr (\((regs, cells), tr) -> insert regs cells tr) empty trs
+-}
+
 instance (PA.ValidArch arch, PSo.ValidSym sym) => W4S.W4Serializable sym (TraceCollection sym arch) where
   w4Serialize (TraceCollection allTraces regs cells) = do
     W4S.object 
diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index 51ee550b..7bccc421 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -1491,19 +1491,23 @@ widenPostcondition scope bundle preD postD0 = do
                Widen widenk (WidenLocs locs) d -> do
                 -- FIXME: should we make a post condition here?
                 tr <- getTraceFromModel scope evalFn bundle preD Nothing
-                let st' = foldr (addLoc widenk tr) prevState locs 
+                let (regs,cells) = getTracedLocs (Set.toList locs)
+                let st' = case widenk of
+                      WidenEquality -> prevState { stTracesEq = PTc.insert regs cells tr (stTracesEq prevState) }
+                      WidenValue -> prevState { stTracesVal = PTc.insert regs cells tr (stTracesVal prevState) }
                 return $ st' { stWidenCase = WidenCaseStep widenk, stLocs = WidenLocs locs <> stLocs prevState, stDomain = d }
                _ -> return $ result res
     where
-      addLoc :: WidenKind -> CE.TraceEvents sym arch -> PL.SomeLocation sym arch ->  WidenState sym arch v -> WidenState sym arch v
-      addLoc wk te (PL.SomeLocation l)  st = case (wk,l) of
-        (WidenEquality, PL.Register r) -> st { stTracesEq = PTc.insertReg r te (stTracesEq st) }
-        (WidenEquality, PL.Cell c) -> st { stTracesEq = PTc.insertCell c te (stTracesEq st) }
-        (WidenValue, PL.Register r) -> st { stTracesVal = PTc.insertReg r te (stTracesVal st) }
-        (WidenValue, PL.Cell c) -> st { stTracesVal = PTc.insertCell c te (stTracesVal st) }
+      getTracedLocs :: [PL.SomeLocation sym arch] -> ([Some (MM.ArchReg arch)], [Some (PMc.MemCell sym arch)])
+      getTracedLocs [] = ([],[])
+      getTracedLocs ((PL.SomeLocation l):locs) =
+        let (regs,cells) = getTracedLocs locs
+        in case l of
+          PL.Register r -> (Some r:regs,cells)
+          PL.Cell c -> (regs,Some c:cells)
         -- other kinds of locations we can ignore for now: since this is just for reporting purposes we
         -- only need to index the traces by locations the user actually knows about
-        _ -> st
+          _ -> (regs,cells)
 
       result :: WidenResult sym arch v -> WidenState sym arch v
       result r = case r of

From 51e489f036a14c5fce5c74cd6bf7a518369f6c5b Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Fri, 11 Oct 2024 11:52:03 -0700
Subject: [PATCH 4/6] add quick solution for using "JSON" tag in repl to view
 JSON for traces collected during widening

---
 src/Pate/Monad.hs                 | 24 ++++++++++++++++++------
 src/Pate/Verification/Widening.hs |  5 ++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/Pate/Monad.hs b/src/Pate/Monad.hs
index 4050b8a4..d69e876b 100644
--- a/src/Pate/Monad.hs
+++ b/src/Pate/Monad.hs
@@ -202,6 +202,8 @@ import Data.List
 import Data.Parameterized.SymbolRepr
 import qualified What4.JSON as W4S
 import qualified Data.Aeson as JSON
+import qualified Data.Aeson.Text as JSON
+import qualified GHC.IO.Unsafe as IO
 
 atPriority :: 
   NodePriority ->
@@ -415,16 +417,16 @@ data W4SerializableNode sym arch nm = IsSerializableNode sym arch nm =>
 
 -- | List of trace nodes (of any type) that share the same expression binding
 --   environment when serialized
-newtype SharedExprEnv sym arch = SharedExprEnv [Some (W4SerializableNode sym arch)]
+data SharedExprEnv sym arch = SharedExprEnv sym [Some (W4SerializableNode sym arch)]
 
 instance ValidSymArch sym arch => IsTraceNode '(sym,arch) "shared_env" where
   type TraceNodeType '(sym,arch) "shared_env" = SharedExprEnv sym arch
   type TraceNodeLabel "shared_env" = String
 
-  prettyNode top_lbl (SharedExprEnv nds) = 
+  prettyNode top_lbl (SharedExprEnv _ nds) =
     PP.vsep $ (PP.pretty top_lbl <> ": "):(map (\(Some (W4SerializableNode (_ :: SymbolRepr nm) lbl nd)) -> prettyNode @_ @'(sym,arch) @nm lbl nd) nds)
 
-  jsonNode sym top_lbl (SharedExprEnv nds) = do
+  jsonNode sym top_lbl (SharedExprEnv _ nds) = do
     (\f -> evalStateT f W4S.emptyExprCache) $ do
       contents <- forM nds $ \(Some (W4SerializableNode nm lbl nd)) -> do
         cache0 <- get
@@ -437,6 +439,16 @@ instance ValidSymArch sym arch => IsTraceNode '(sym,arch) "shared_env" where
       env_json <- liftIO $ W4S.serializeExprEnv sym env
       return $ JSON.object [ "name" JSON..= top_lbl, "shared_env" JSON..= env_json, "contents" JSON..= contents]
 
+  nodeTags = [(Summary, prettyNode @_ @'(sym,arch) @"shared_env")
+              -- fixme: quick hack to get JSON out without needing to run the entire
+              -- toplevel in JSON mode. Ideally this would be something that the Repl could
+              -- do itself.
+              , ("JSON", \top_lbl env@(SharedExprEnv sym _) ->
+                  let v = IO.unsafePerformIO $ jsonNode @_ @'(sym,arch) @"shared_env" sym top_lbl env
+                  in PP.pretty $ JSON.encodeToLazyText v)
+             ]
+
+
 -- | Emit multiple values to the tracetree with a common expression binding environment
 --   This will ultimately result in a single composite value, where all of the given values are serialized
 --   using the same expression binding cache.
@@ -456,12 +468,12 @@ withSharedEnvEmit ::
   String ->
   (forall m. Monad m => (forall nm. (IsSerializableNode sym arch nm) => SymbolRepr nm -> TraceNodeLabel nm -> TraceNodeType '(sym,arch) nm -> m ()) -> m ()) ->
   EquivM sym arch ()
-withSharedEnvEmit top_nm f = do
-  env <- execStateT (f g) (SharedExprEnv [])
+withSharedEnvEmit top_nm f = withSym $ \sym -> do
+  env <- execStateT (f g) (SharedExprEnv sym [])
   emitTraceLabel @"shared_env" @'(sym,arch) top_nm env
   where 
     g :: forall (nm :: Symbol). IsSerializableNode sym arch nm => SymbolRepr nm -> TraceNodeLabel nm -> TraceNodeType '(sym,arch) nm -> StateT (SharedExprEnv sym arch) (EquivM_ sym arch) ()
-    g nm lbl v = modify $ \(SharedExprEnv vs) -> SharedExprEnv ((Some (W4SerializableNode nm lbl v)):vs)
+    g nm lbl v = modify $ \(SharedExprEnv sym vs) -> SharedExprEnv sym ((Some (W4SerializableNode nm lbl v)):vs)
 
 withBinary ::
   forall bin sym arch a.
diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index 7bccc421..4078b125 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -1400,11 +1400,14 @@ widenPostcondition ::
 widenPostcondition scope bundle preD postD0 = do
   st <- withTracing @"debug" "widenPostcondition" $ withSym $ \sym -> do
     eqCtx <- equivalenceContext
+    initSt <- initWidenState bundle postD0
     traceBundle bundle "Entering widening loop"
     subTree @"domain" "Widening Steps" $
-      widenLoop sym localWideningGas eqCtx (initWidenState postD0)
+      widenLoop sym localWideningGas eqCtx initSt
     
   case stWidenCase st of
+    -- we use 'withSharedEnvEmit' so that the underlying 'TraceCollection's are serialized
+    -- in a shared environment with the domain.
     WidenCaseStep _ -> withSharedEnvEmit "Equivalence Counter-example Traces" $ \emit -> do
       emit (CT.knownSymbol @"trace_collection") "Equality" (stTracesEq st)
       emit (CT.knownSymbol @"trace_collection") "Value" (stTracesVal st)

From 2fbde9e027cae5a99756b8078a0e07c913bbbbd3 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Fri, 11 Oct 2024 11:53:03 -0700
Subject: [PATCH 5/6] tryWithAsms: check satisfiability first to optimize for
 unsat case

this avoids redundant satisfiability checks for assumptions,
which otherwise causes a significant slowdown when used during each
widening step
---
 src/Pate/Verification/Widening.hs | 119 +++++++++++++++++++++---------
 1 file changed, 85 insertions(+), 34 deletions(-)

diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index 4078b125..6f5ad3ba 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -481,27 +481,62 @@ addRefinementChoice nd gr0 = withTracing @"message" ("Modify Proof Node: " ++ sh
       choice "Clear work list (unsafe!)" $ \_ gr4 ->
         return $ emptyWorkList gr4
 
-tryWithAsms :: 
-  [(W4.Pred sym, PEE.InnerEquivalenceError arch)] ->
+-- | Predicates which we attempt to assume when generating a model, but raise
+--   the corresponding warning when this is not possible.
+type WeakAssumptions sym arch = [(W4.Pred sym, Maybe (PEE.InnerEquivalenceError arch))]
+
+-- | Given a satisfiable predicate 'p', call the given continuation
+--   with a model where a maximal subset of the given assumptions are also
+--   satisfied.
+--   Throws the given error for each assumption that is not satisfied in the model, and
+--   ultimately returns 'Nothing' if 'p' is unsatisfiable (or unknown) on its own.
+
+-- | Check the given predicate 'p' for satisfiability, while computing a model
+--   where as many of the given assumptions are true as possible (when 'p' is satisfiable).
+tryWithAsms ::
+  forall sym arch a.
+  WeakAssumptions sym arch ->
   W4.Pred sym ->
   (SatResult (SymGroundEvalFn sym) () -> EquivM_ sym arch a) ->
   EquivM sym arch a
-tryWithAsms ((asm, err):asms) p f = do
-  mresult <- withSatAssumption (PAs.fromPred asm) $ tryWithAsms asms p f
+tryWithAsms asms_ p f = do
+  -- NOTE: we avoid 'withSatAssumption' here since that's somewhat more
+  -- expensive than just 'goalSat', and we'd prefer to fail quickly
+  mresult <- goalSat "check" p $ \case
+    Unsat () -> Just <$> f (Unsat ())
+    Unknown -> Just <$> f Unknown
+    -- ideally we could just call 'go asms_' from here directly,
+    -- but nested solver calls seem to break What4, so there's a
+    -- bit of redundancy here
+    Sat{} -> return Nothing
   case mresult of
     Just a -> return a
-    Nothing -> emitWarning err >> tryWithAsms asms p f
-tryWithAsms [] p f = goalSat "tryWithAsms" p f
-
--- | Check the predicate for satisfiability, attempting to
---   constrain the model to valid pointers and zero-offset stacks
---   if possible.
-withTraceAssumptions ::
-  SimBundle sym arch v ->
-  W4.Pred sym ->
-  (SatResult (SymGroundEvalFn sym) () -> EquivM_ sym arch a) ->
-  EquivM sym arch a
-withTraceAssumptions bundle p f = withSym $ \sym -> do
+    Nothing -> withAssumption p $ go asms_
+  where
+    maybeEmit :: Maybe (PEE.InnerEquivalenceError arch) -> EquivM_ sym arch ()
+    maybeEmit Nothing = return ()
+    maybeEmit (Just err) = emitWarning err
+
+    go ::
+      [(W4.Pred sym, Maybe (PEE.InnerEquivalenceError arch))] ->
+      EquivM_ sym arch a
+    go [] = withSym $ \sym -> goalSat "go" (W4.truePred sym) f
+    go [(asm, err)] = do
+      mresult <- goalSat "go" asm $ \case
+        Sat evalFn -> Just <$> f (Sat evalFn)
+        Unsat () -> return Nothing
+        Unknown -> return Nothing
+      case mresult of
+        Nothing -> maybeEmit err >> go []
+        Just a -> return a
+    go ((asm, err):asms) = do
+      mresult <- withSatAssumption (PAs.fromPred asm) $ go asms
+      case mresult of
+        Nothing -> maybeEmit err >> go asms
+        Just a -> return a
+
+mkTraceAssumptions :: SimBundle sym arch v -> EquivM sym arch (WeakAssumptions sym arch)
+mkTraceAssumptions bundle = withSym $ \sym -> do
   (_, ptrAsserts) <- PVV.collectPointerAssertions bundle
   stacks_zero <- PPa.catBins $ \bin -> do
     in_ <- PPa.get bin (PS.simIn bundle)
@@ -509,11 +544,12 @@ withTraceAssumptions bundle p f = withSym $ \sym -> do
     zero <- liftIO $ W4.bvLit sym CT.knownRepr (BVS.mkBV CT.knownRepr 0)
     PAs.fromPred <$> (liftIO $ W4.isEq sym zero stackbase)
 
-  ptrAsserts_pred <- PAs.toPred sym (stacks_zero <> ptrAsserts)
-
-  tryWithAsms
-    [ (ptrAsserts_pred, PEE.RequiresInvalidPointerOps)
-    ] p f
+  stacks_zero_pred <- PAs.toPred sym stacks_zero
+  ptrAsserts_pred <- PAs.toPred sym ptrAsserts
+  return $
+    [ (stacks_zero_pred, Nothing )
+    , (ptrAsserts_pred, Just PEE.RequiresInvalidPointerOps)
+    ]
 
 getSomeGroundTrace ::
   PS.SimScope sym arch v ->
@@ -522,7 +558,8 @@ getSomeGroundTrace ::
   Maybe (StatePostCondition sym arch v) ->
   EquivM sym arch (CE.TraceEvents sym arch)
 getSomeGroundTrace scope bundle preD postCond = withSym $ \sym -> do
-  withTraceAssumptions bundle (W4.truePred sym) $ \case
+  trace_asms <- mkTraceAssumptions bundle
+  tryWithAsms trace_asms (W4.truePred sym) $ \case
     Unsat{} -> throwHere PEE.InvalidSMTModel
     Unknown -> throwHere PEE.InconclusiveSAT
     Sat evalFn -> getTraceFromModel scope evalFn bundle preD postCond
@@ -1356,11 +1393,18 @@ data WidenState sym arch v = WidenState
   , stTracesEq :: PTc.TraceCollection sym arch
     -- ^ collected traces for equality widening steps
   , stTracesVal :: PTc.TraceCollection sym arch
-    --- ^ collected traces for value widening steps
+    -- ^ collected traces for value widening steps
+  , stTraceAsms :: WeakAssumptions sym arch
+    -- ^ pre-computed assumptions to try to add when generating traces (i.e. from withTraceAssumptions)
   }
 
-initWidenState :: AbstractDomain sym arch v -> WidenState sym arch v
-initWidenState d = WidenState d (WidenLocs Set.empty) WidenCaseStart PTc.empty PTc.empty
+initWidenState ::
+  SimBundle sym arch v ->
+  AbstractDomain sym arch v ->
+  EquivM sym arch (WidenState sym arch v)
+initWidenState bundle d = do
+  traceAsms <- mkTraceAssumptions bundle
+  return $ WidenState d (WidenLocs Set.empty) WidenCaseStart PTc.empty PTc.empty traceAsms
 
 
 -- Compute a final 'WidenResult' from an (intermediate) 'WidenState' (surjective)
@@ -1425,21 +1469,26 @@ widenPostcondition scope bundle preD postD0 = do
      EquivM_ sym arch (WidenState sym arch v)
    widenOnce widenK (Gas i) mwb prevState loc goal = case stWidenCase prevState of
      WidenCaseErr{} -> return prevState
-     _ -> startTimer $ withSym $ \sym -> do
+     _ -> withSym $ \sym -> do
        eqCtx <- equivalenceContext
        let 
         this_loc = WidenLocs (Set.singleton (PL.SomeLocation loc))
         postD = stDomain prevState
+
+       {-
        curAsms <- currentAsm
        let emit r =
              withValid @() $ emitEvent (PE.SolverEvent (PS.simPair bundle) PE.EquivalenceProof r curAsms goal)
        emit PE.SolverStarted
+       -}
+       let emit _r = return ()
+
 
        not_goal <- liftIO $ W4.notPred sym goal
        
        --(not_goal', ptrAsms) <- PVV.collectPointerAssertions not_goal
        emitTraceLabel @"expr" "goal" (Some goal)
-       withTraceAssumptions bundle not_goal $ \case
+       tryWithAsms (stTraceAsms prevState) not_goal $ \case
          Unsat _ -> do
            emit PE.SolverSuccess
            return prevState
@@ -1492,13 +1541,15 @@ widenPostcondition scope bundle preD postD0 = do
                    return $ result $ WideningError msg this_loc postD
                  _ -> return prevState
                Widen widenk (WidenLocs locs) d -> do
-                -- FIXME: should we make a post condition here?
-                tr <- getTraceFromModel scope evalFn bundle preD Nothing
-                let (regs,cells) = getTracedLocs (Set.toList locs)
-                let st' = case widenk of
-                      WidenEquality -> prevState { stTracesEq = PTc.insert regs cells tr (stTracesEq prevState) }
-                      WidenValue -> prevState { stTracesVal = PTc.insert regs cells tr (stTracesVal prevState) }
-                return $ st' { stWidenCase = WidenCaseStep widenk, stLocs = WidenLocs locs <> stLocs prevState, stDomain = d }
+                 let nextState = prevState
+                      { stWidenCase = WidenCaseStep widenk
+                      , stLocs = WidenLocs locs <> stLocs prevState
+                      , stDomain = d }
+                 tr <- getTraceFromModel scope evalFn bundle preD Nothing
+                 let (regs,cells) = getTracedLocs (Set.toList locs)
+                 return $ case widenk of
+                   WidenEquality -> nextState { stTracesEq = PTc.insert regs cells tr (stTracesEq nextState) }
+                   WidenValue -> nextState { stTracesVal = PTc.insert regs cells tr (stTracesVal nextState) }
                _ -> return $ result res
     where
       getTracedLocs :: [PL.SomeLocation sym arch] -> ([Some (MM.ArchReg arch)], [Some (PMc.MemCell sym arch)])

From 1cb3761b44b8940fc00ecef54846fb284932c775 Mon Sep 17 00:00:00 2001
From: Daniel Matichuk <dmatichuk@galois.com>
Date: Mon, 14 Oct 2024 09:41:47 -0700
Subject: [PATCH 6/6] tryWithAsms: documentation cleanup

---
 src/Pate/Verification/Widening.hs | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/src/Pate/Verification/Widening.hs b/src/Pate/Verification/Widening.hs
index 6f5ad3ba..800df75d 100644
--- a/src/Pate/Verification/Widening.hs
+++ b/src/Pate/Verification/Widening.hs
@@ -485,14 +485,11 @@ addRefinementChoice nd gr0 = withTracing @"message" ("Modify Proof Node: " ++ sh
 --   the corresponding warning when this is not possible.
 type WeakAssumptions sym arch = [(W4.Pred sym, Maybe (PEE.InnerEquivalenceError arch))]
 
--- | Given a satisfiable predicate 'p', call the given continuation
---   with a model where a maximal subset of the given assumptions are also
+-- | Given a predicate 'p', call the given continuation 'f'
+--   with a model where a maximal subset of the given 'WeakAssumptions' are also
 --   satisfied.
---   Throws the given error for each assumption that is not satisfied in the model, and
---   ultimately returns 'Nothing' if 'p' is unsatisfiable (or unknown) on its own.
-
--- | Check the given predicate 'p' for satisfiability, while computing a model
---   where as many of the given assumptions are true as possible (when 'p' is satisfiable).
+--   If 'p' is not satisfiable (or inconclusive) on its own then the given assumptions are ignored,
+--   and 'f' is called with the result of checking just 'p' (either 'Unsat' or 'Unknown').
 tryWithAsms ::
   forall sym arch a.
   WeakAssumptions sym arch ->
@@ -1475,14 +1472,10 @@ widenPostcondition scope bundle preD postD0 = do
         this_loc = WidenLocs (Set.singleton (PL.SomeLocation loc))
         postD = stDomain prevState
 
-       {-
        curAsms <- currentAsm
        let emit r =
              withValid @() $ emitEvent (PE.SolverEvent (PS.simPair bundle) PE.EquivalenceProof r curAsms goal)
        emit PE.SolverStarted
-       -}
-       let emit _r = return ()
-
 
        not_goal <- liftIO $ W4.notPred sym goal