GaloisInc · travitch · Jan 14, 2021
diff --git a/surveyor-brick/src/Surveyor/Brick/Widget/ValueViewer.hs b/surveyor-brick/src/Surveyor/Brick/Widget/ValueViewer.hs
@@ -27,10 +27,7 @@ import qualified Control.Lens as L
 import qualified Control.Monad.State.Strict as St
 import qualified Data.BitVector.Sized as DBS
 import qualified Data.List as L
-import qualified Data.List.NonEmpty as DLN
-import qualified Data.Map.Strict as Map
 import           Data.Maybe ( fromMaybe )
-import qualified Data.Parameterized.Map as MapF
 import qualified Data.Parameterized.Nonce as PN
 import           Data.Parameterized.Some ( Some(..) )
 import qualified Data.Parameterized.TraversableFC as FC
@@ -55,6 +52,7 @@ import qualified What4.Utils.StringLiteral as WUS
 
 import           Surveyor.Brick.Names ( Names(..) )
 import qualified Surveyor.Core as C
+import qualified Surveyor.Core.ExprMap as SCEM
 
 -- | Mark widgets as either being rendered inline in operand positions or as
 -- generating a binding that will be referred to by a number in an operand
@@ -77,23 +75,6 @@ data RenderWidget sym =
                , rwValue :: Maybe (Some (LMCR.RegEntry sym))
                }
 
--- | This is a simple wrapper around a 'RenderWidget' to throw away a type
--- parameter so that we can store it in a 'MapF.MapF'.  We need a custom type
--- instead of just using 'Data.Functor.Const' because that type is too
--- polymorphic and causes some type inference problems.
-newtype ConstWidget sym (tp :: WT.BaseType) = ConstWidget (RenderWidget sym)
-
--- This is a very simple wrapper around the 'WEB.Expr' type arranged so that we
--- can use them as surrogate keys in the cache.  The challenge is that we can't
--- expose the ExprBuilder in the type (without propagating it everywhere), but
--- we need it in order to get an Ord instance in scope.  Instead, we capture
--- that equality in an existential, which is enough to let us derive Eq/Ord.
-data SymExpr s sym where
-  SymExpr :: (sym ~ WEB.ExprBuilder s st fs) => Some (WEB.Expr s) -> SymExpr s sym
-
-deriving instance Eq (SymExpr s sym)
-deriving instance Ord (SymExpr s sym)
-
 -- | A cache of pre-rendered widgets
 --
 -- Re-rendering some widgets is costly (and we need to avoid it to preserve
@@ -104,13 +85,12 @@ deriving instance Ord (SymExpr s sym)
 -- caches could be unified on a sum type for keys, but keeping them separate is
 -- more efficient for the common case (of nonce lookups)
 data Cache s sym =
-  Cache { _cachedByNonce :: !(MapF.MapF (PN.Nonce s) (ConstWidget sym))
-        , _cachedBySurrogate :: !(Map.Map (DLN.NonEmpty (SymExpr s sym)) (RenderWidget sym))
+  Cache { _cachedWidgets :: !(SCEM.ExprMap s sym (RenderWidget sym))
         , _surrogateIdentifiers :: !Int
         }
 
 emptyCache :: Cache s sym
-emptyCache = Cache MapF.empty mempty 0
+emptyCache = Cache SCEM.emptyExprMap 0
 
 L.makeLenses ''Cache
 
@@ -198,14 +178,11 @@ buildViewer tp re = do
   -- We'll use this same primitive for rendering later, but then we'll start it
   -- with a primed cache
   root <- buildTermWidget tp re
-  nonceEntries <- St.gets (^. cachedByNonce)
-  surrogateEntries <- St.gets (^. cachedBySurrogate)
+  cached <- St.gets (^. cachedWidgets)
   c <- St.get
   -- FIXME: It would be really nice to sort these in dependency order; being
   -- split between two caches, entries could be interleaved oddly
-  let cachedVals = concat [ [ w | MapF.Pair _ (ConstWidget w) <- MapF.toList nonceEntries ]
-                          , [ w | (_, w) <- Map.toList surrogateEntries ]
-                          ]
+  let cachedVals = SCEM.elems cached
   let vals = DV.fromList (if null cachedVals then [root] else cachedVals)
   let l0 = BL.list ValueViewerList vals 1
   let vvs = ValueViewerState { regType = tp
@@ -227,19 +204,19 @@ buildTermWidget tp re =
         LCT.UnitRepr ->
           -- We don't update the cache here because we don't have a nonce for these
           inline (B.txt "()")
-        CLM.LLVMPointerRepr _w ->
-          case re of
-            CLM.LLVMPointer base off -> do
-              let key = SymExpr (Some base) DLN.:| [SymExpr (Some off)]
-              surrogateCache <- St.gets (^. cachedBySurrogate)
-              case Map.lookup key surrogateCache of
-                Just rw -> return rw
-                Nothing -> do
+        CLM.LLVMPointerRepr _w -> do
+          let key = SCEM.llvmPointerKey (LMCR.RegEntry tp re)
+          cached <- St.gets (^. cachedWidgets)
+          case SCEM.lookupExprValue key cached of
+            Just rw -> return rw
+            Nothing -> do
+              case re of
+                CLM.LLVMPointer base off -> do
                   base' <- argRef <$> buildTermWidget (LCT.baseToType (WI.exprType base)) base
                   off' <- argRef <$> buildTermWidget (LCT.baseToType (WI.exprType off)) off
                   let rendering = intersperse [B.txt "llvmPointer", base', off']
                   bindSurrogate key (LMCR.RegEntry tp re) rendering
-                  -- inline (intersperse [B.txt "llvmPointer", base', off'])
+
         _ -> inline (B.txt ("Unhandled crucible type " <> T.pack (show tp)))
     LCT.AsBaseType _btp ->
       case re of
@@ -256,9 +233,10 @@ buildTermWidget tp re =
           inlineLoc re (B.txt "$" B.<+> B.txt (WS.solverSymbolAsText (WEB.bvarName bv)))
         WEB.NonceAppExpr nae -> do
           let nonce = WEB.nonceExprId nae
-          m <- St.gets (^. cachedByNonce)
-          case MapF.lookup nonce m of
-            Just (ConstWidget cached) -> return cached
+          let key = SCEM.nonceValueKey nonce
+          cached <- St.gets (^. cachedWidgets)
+          case SCEM.lookupExprValue key cached of
+            Just cachedRendering -> return cachedRendering
             Nothing ->
               case WEB.nonceExprApp nae of
                 WEB.Annotation bt annotNonce e -> do
@@ -306,9 +284,10 @@ buildTermWidget tp re =
                                             ])
         WEB.AppExpr ae -> do
           let nonce = WEB.appExprId ae
-          m <- St.gets (^. cachedByNonce)
-          case MapF.lookup nonce m of
-            Just (ConstWidget cached) -> return cached
+          let key = SCEM.nonceValueKey nonce
+          cached <- St.gets (^. cachedWidgets)
+          case SCEM.lookupExprValue key cached of
+            Just cachedRendering -> return cachedRendering
             Nothing ->
               case WEB.appExprApp ae of
                 WEB.BaseEq _tp e1 e2 -> bindBinExpr ae e1 e2 "eq"
@@ -566,12 +545,13 @@ bindExpr :: (HasNonce e, sym ~ WEB.ExprBuilder t st fs)
          -> B.Widget Names
          -> ViewerBuilder t sym (RenderWidget sym)
 bindExpr ae w =
-  returnCached (getNonce ae) $! RenderWidget (RenderBound (thisRef ae) w) (Just (getLoc ae)) (Just regEntry)
+  returnCached key $! RenderWidget (RenderBound (thisRef ae) w) (Just (getLoc ae)) (Just regEntry)
   where
+    key = SCEM.nonceValueKey (getNonce ae)
     re = getSymExpr ae
     regEntry = Some (LMCR.RegEntry (LCT.baseToType (WI.exprType re)) re)
 
-bindSurrogate :: DLN.NonEmpty (SymExpr s sym)
+bindSurrogate :: SCEM.ExprMapKey s sym tp
               -> LMCR.RegEntry sym tp
               -> B.Widget Names
               -> ViewerBuilder s sym (RenderWidget sym)
@@ -582,15 +562,15 @@ bindSurrogate key entry w = do
   -- different namespace
   let binding = B.str (printf "%%%d" nextSurrogate)
   let rw = RenderWidget (RenderBound binding w) Nothing (Just (Some entry))
-  St.modify' $ \s -> s & cachedBySurrogate %~ Map.insert key rw
+  St.modify' $ \s -> s & cachedWidgets %~ SCEM.addExprValue key rw
   return rw
 
-returnCached :: forall s sym (tp :: WT.BaseType)
-              . PN.Nonce s tp
+returnCached :: forall s sym tp
+              . SCEM.ExprMapKey s sym tp
              -> RenderWidget sym
              -> ViewerBuilder s sym (RenderWidget sym)
-returnCached nonce res = do
-  St.modify' $ \s -> s & cachedByNonce %~ MapF.insert nonce (ConstWidget res)
+returnCached key res = do
+  St.modify' $ \s -> s & cachedWidgets %~ SCEM.addExprValue key res
   return res
 
 thisRef :: (HasNonce e) => e t tp -> B.Widget n

diff --git a/surveyor-core/src/Surveyor/Core/ExprMap.hs b/surveyor-core/src/Surveyor/Core/ExprMap.hs
@@ -0,0 +1,101 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE StandaloneDeriving #-}
+-- | This module defines helpers for mapping from symbolic values, specifically
+-- 'RegEntry', to arbitrary values.
+--
+-- This is a challenge because there is no Ord instance for 'RegEntry', nor
+-- nonces.  Some subset of values have 'Ord' instances, so we decompose
+-- composite values that do not into surrogate keys (in terms of 'SymExpr')
+module Surveyor.Core.ExprMap (
+  SymExpr(..),
+  ExprMap,
+  emptyExprMap,
+  addExprValue,
+  lookupExprValue,
+  elems,
+  ExprMapKey,
+  baseValueKey,
+  nonceValueKey,
+  llvmPointerKey
+  ) where
+
+import qualified Data.List.NonEmpty as DLN
+import qualified Data.Map.Strict as Map
+import qualified Data.Parameterized.Map as MapF
+import qualified Data.Parameterized.Nonce as PN
+import           Data.Parameterized.Some ( Some(..) )
+import qualified Lang.Crucible.LLVM.MemModel as CLM
+import qualified Lang.Crucible.Simulator.RegMap as LMCR
+import qualified Lang.Crucible.Types as LCT
+import qualified What4.BaseTypes as WT
+import qualified What4.Expr.Builder as WEB
+
+-- | A wrapper around values of type @a sym@ that provides a phantom base type parameter
+newtype BaseTypeValue a (tp :: WT.BaseType) = BaseTypeValue { unwrapBaseTypeValue :: a }
+
+-- This is a very simple wrapper around the 'WEB.Expr' type arranged so that we
+-- can use them as surrogate keys in the cache.  The challenge is that we can't
+-- expose the ExprBuilder in the type (without propagating it everywhere), but
+-- we need it in order to get an Ord instance in scope.  Instead, we capture
+-- that equality in an existential, which is enough to let us derive Eq/Ord.
+data SymExpr s sym where
+  SymExpr :: (sym ~ WEB.ExprBuilder s st fs) => Some (WEB.Expr s) -> SymExpr s sym
+
+deriving instance Eq (SymExpr s sym)
+deriving instance Ord (SymExpr s sym)
+
+data ExprMapKey s sym (tp :: LCT.CrucibleType) where
+  NonceKey :: !(PN.Nonce s bt) -> ExprMapKey s sym (LCT.BaseToType bt)
+  SurrogateKey :: !(DLN.NonEmpty (SymExpr s sym)) -> ExprMapKey s sym tp
+
+data ExprMap s sym a =
+  ExprMap { nonceMap :: !(MapF.MapF (PN.Nonce s) (BaseTypeValue a))
+          , surrogateMap :: !(Map.Map (DLN.NonEmpty (SymExpr s sym)) a)
+          }
+
+emptyExprMap :: ExprMap s sym a
+emptyExprMap = ExprMap { nonceMap = MapF.empty, surrogateMap = mempty }
+
+llvmPointerKey :: (sym ~ WEB.ExprBuilder s st fs)
+               => LMCR.RegEntry sym (CLM.LLVMPointerType w)
+               -> ExprMapKey s sym (CLM.LLVMPointerType w)
+llvmPointerKey regEntry =
+  case LMCR.regValue regEntry of
+    CLM.LLVMPointer base off -> SurrogateKey ((SymExpr (Some base)) DLN.:| [SymExpr (Some off)])
+
+baseValueKey :: (sym ~ WEB.ExprBuilder s st fs)
+             => LMCR.RegEntry sym (LCT.BaseToType bt)
+             -> Maybe (ExprMapKey s sym (LCT.BaseToType bt))
+baseValueKey regEntry = NonceKey <$> WEB.exprMaybeId (LMCR.regValue regEntry)
+
+nonceValueKey :: ( sym ~ WEB.ExprBuilder s st fs
+                 , tp ~ LCT.BaseToType bt
+                 )
+              => PN.Nonce s bt
+              -> ExprMapKey s sym tp
+nonceValueKey nonce = NonceKey nonce
+
+-- | Add a value to the 'ExprMap' at the given key
+addExprValue :: ExprMapKey s sym tp
+             -> a
+             -> ExprMap s sym a
+             -> ExprMap s sym a
+addExprValue key val m =
+  case key of
+    NonceKey nonce -> m { nonceMap = MapF.insert nonce (BaseTypeValue val) (nonceMap m) }
+    SurrogateKey sk -> m { surrogateMap = Map.insert sk val (surrogateMap m) }
+
+lookupExprValue :: ExprMapKey s sym tp
+                -> ExprMap s sym a
+                -> Maybe a
+lookupExprValue key m =
+  case key of
+    NonceKey nonce -> unwrapBaseTypeValue <$> MapF.lookup nonce (nonceMap m)
+    SurrogateKey sk -> Map.lookup sk (surrogateMap m)
+
+elems :: ExprMap s sym a -> [a]
+elems (ExprMap nm skm) = Map.elems skm ++ [ v
+                                          | Some (BaseTypeValue v) <- MapF.elems nm
+                                          ]
diff --git a/surveyor-core/src/Surveyor/Core/GraphViz.hs b/surveyor-core/src/Surveyor/Core/GraphViz.hs
@@ -15,8 +15,7 @@ import qualified Data.BitVector.Sized as DBS
 import qualified Data.Foldable as F
 import qualified Data.GraphViz as DG
 import qualified Data.Map.Strict as Map
-import           Data.Maybe ( fromMaybe, isJust )
-import qualified Data.Parameterized.Classes as PC
+import           Data.Maybe ( fromMaybe )
 import qualified Data.Parameterized.Context as Ctx
 import qualified Data.Parameterized.Nonce as PN
 import qualified Data.Parameterized.TraversableFC as FC
@@ -35,6 +34,8 @@ import qualified What4.Symbol as WS
 import qualified What4.Utils.Complex as WUC
 import qualified What4.Utils.StringLiteral as WUS
 
+import qualified Surveyor.Core.ExprMap as SCEM
+
 -- | How a sub-term should be rendered in its host term
 --
 -- This will either be a unique identifier @$n@ or an inlined constant
@@ -88,26 +89,17 @@ data FormulaEdgeLabel where
 deriving instance Eq FormulaEdgeLabel
 deriving instance Ord FormulaEdgeLabel
 
-data CrucibleNonce s where
-  CrucibleNonce :: forall s (tp :: LCT.BaseType) . PN.Nonce s tp -> CrucibleNonce s
-
-instance Eq (CrucibleNonce s) where
-  CrucibleNonce n1 == CrucibleNonce n2 = isJust (PC.testEquality n1 n2)
-
-instance Ord (CrucibleNonce s) where
-  compare (CrucibleNonce n1) (CrucibleNonce n2) = PC.toOrdering (PC.compareF n1 n2)
-
 data RegEntryBuilderState s sym =
   RegEntryBuilderState { builderEdges :: Map.Map (FormulaNode s sym) (Set.Set (FormulaNode s sym, FormulaEdgeLabel))
                        , builderNodes :: Map.Map (FormulaNode s sym) FormulaNodeLabel
-                       , renderCache :: Map.Map (CrucibleNonce s) Rendering
+                       , renderCache :: SCEM.ExprMap s sym Rendering
                        }
 
 emptyRegEntryBuilderState :: RegEntryBuilderState s sym
 emptyRegEntryBuilderState =
   RegEntryBuilderState { builderEdges = Map.empty
                        , builderNodes = Map.empty
-                       , renderCache = Map.empty
+                       , renderCache = SCEM.emptyExprMap
                        }
 
 newtype RegEntryBuilder s sym a =
@@ -167,7 +159,7 @@ traverseFormulaStructure predEntry entry =
           let regVal = LMCR.regValue entry
           addEdge predEntry n
           cache <- MS.gets renderCache
-          case Map.lookup (CrucibleNonce nonce) cache of
+          case SCEM.lookupExprValue (SCEM.nonceValueKey nonce) cache of
             Just r -> return r
             Nothing ->
               case WEB.nonceExprApp nae of
@@ -223,7 +215,7 @@ traverseFormulaStructure predEntry entry =
           let regVal = LMCR.regValue entry
           addEdge predEntry n
           cache <- MS.gets renderCache
-          case Map.lookup (CrucibleNonce nonce) cache of
+          case SCEM.lookupExprValue (SCEM.nonceValueKey nonce) cache of
             Just r -> return r
             Nothing ->
               case WEB.appExprApp ae of
@@ -466,16 +458,18 @@ ensureNode n nl = MS.modify' $ \g -> g { builderEdges = Map.insertWith Set.union
                                        , builderNodes = Map.insert n nl (builderNodes g)
                                        }
 
-bind :: forall sym s st fs (tp :: LCT.BaseType)
-      . (sym ~ WEB.ExprBuilder s st fs)
-     => PN.Nonce s tp
-     -> WEB.Expr s tp
+bind :: forall sym s st fs (bt :: LCT.BaseType) tp
+      . ( sym ~ WEB.ExprBuilder s st fs
+        , tp ~ LCT.BaseToType bt
+        )
+     => PN.Nonce s bt
+     -> WEB.Expr s bt
      -> PP.Doc ()
      -> RegEntryBuilder s sym Rendering
 bind nonce e doc = do
   MS.modify' $ \g -> g { builderEdges = Map.insertWith Set.union n Set.empty (builderEdges g)
                        , builderNodes = Map.insert n nl (builderNodes g)
-                       , renderCache = Map.insert (CrucibleNonce nonce) ref (renderCache g)
+                       , renderCache = SCEM.addExprValue (SCEM.nonceValueKey nonce) ref (renderCache g)
                        }
   return ref
   where

diff --git a/surveyor-core/surveyor-core.cabal b/surveyor-core/surveyor-core.cabal
@@ -16,6 +16,7 @@ library
   exposed-modules:     Control.NF
                        Control.Once
                        Surveyor.Core
+                       Surveyor.Core.ExprMap
   other-modules:       Surveyor.Core.Architecture
                        Surveyor.Core.Architecture.Class
                        Surveyor.Core.Architecture.Crucible