Rework memoMap to squash space leaks.

The main benefit of this reorganization is that it
notices when a memoized `SeqMap` has been forced
at all of its locations.  This allows us to
discard the underlying computation, which will
never need to be consulted again.  This, in turn,
allows the garbage collector to reclaim the associated
memory and help prevent certain classes of space leaks.

src/Cryptol/Backend/SeqMap.hs

@@ -62,24 +62,43 @@ import Cryptol.Backend.Concrete (Concrete)
 import Cryptol.Backend.Monad (Unsupported(..))
 
 import Cryptol.TypeCheck.Solver.InfNat(Nat'(..))
+import Cryptol.Utils.Panic
 
 -- | A sequence map represents a mapping from nonnegative integer indices
 --   to values.  These are used to represent both finite and infinite sequences.
 data SeqMap sym a
   = IndexSeqMap  !(Integer -> SEval sym a)
   | UpdateSeqMap !(Map Integer (SEval sym a))
-                 !(Integer -> SEval sym a)
-
+                 !(SeqMap sym a)
+  | MemoSeqMap
+     !Nat'
+     !(IORef (Map Integer a))
+     !(IORef (Integer -> SEval sym a))
 
 indexSeqMap :: (Integer -> SEval sym a) -> SeqMap sym a
 indexSeqMap = IndexSeqMap
 
-lookupSeqMap :: SeqMap sym a -> Integer -> SEval sym a
+lookupSeqMap :: Backend sym => SeqMap sym a -> Integer -> SEval sym a
 lookupSeqMap (IndexSeqMap f) i = f i
-lookupSeqMap (UpdateSeqMap m f) i =
+lookupSeqMap (UpdateSeqMap m xs) i =
   case Map.lookup i m of
     Just x  -> x
-    Nothing -> f i
+    Nothing -> lookupSeqMap xs i
+lookupSeqMap (MemoSeqMap sz cache eval) i =
+  do mz <- liftIO (Map.lookup i <$> readIORef cache)
+     case mz of
+       Just z  -> return z
+       Nothing ->
+         do f <- liftIO (readIORef eval)
+            v <- f i
+            msz <- liftIO $ atomicModifyIORef' cache (\m ->
+                        let m' = Map.insert i v m in (m', Map.size m'))
+            -- If we memoize the entire map, overwrite the evaluation closure to let
+            -- the garbage collector reap it
+            when (case sz of Inf -> False; Nat sz' -> toInteger msz >= sz')
+                 (liftIO (writeIORef eval
+                    (\j -> panic "lookupSeqMap" ["Messed up size accounting", show sz, show j])))
+            return v
 
 instance Backend sym => Functor (SeqMap sym) where
   fmap f xs = IndexSeqMap (\i -> f <$> lookupSeqMap xs i)
@@ -89,36 +108,37 @@ finiteSeqMap :: Backend sym => sym -> [SEval sym a] -> SeqMap sym a
 finiteSeqMap sym xs =
    UpdateSeqMap
       (Map.fromList (zip [0..] xs))
-      (\i -> invalidIndex sym i)
+      (IndexSeqMap (\i -> invalidIndex sym i))
 
 -- | Generate an infinite sequence map from a stream of values
 infiniteSeqMap :: Backend sym => sym -> [SEval sym a] -> SEval sym (SeqMap sym a)
 infiniteSeqMap sym xs =
    -- TODO: use an int-trie?
-   memoMap sym (IndexSeqMap $ \i -> genericIndex xs i)
+   memoMap sym Inf (IndexSeqMap $ \i -> genericIndex xs i)
 
 -- | Create a finite list of length @n@ of the values from @[0..n-1]@ in
 --   the given the sequence emap.
-enumerateSeqMap :: (Integral n) => n -> SeqMap sym a -> [SEval sym a]
+enumerateSeqMap :: (Backend sym, Integral n) => n -> SeqMap sym a -> [SEval sym a]
 enumerateSeqMap n m = [ lookupSeqMap m  i | i <- [0 .. (toInteger n)-1] ]
 
 -- | Create an infinite stream of all the values in a sequence map
-streamSeqMap :: SeqMap sym a -> [SEval sym a]
+streamSeqMap :: Backend sym => SeqMap sym a -> [SEval sym a]
 streamSeqMap m = [ lookupSeqMap m i | i <- [0..] ]
 
 -- | Reverse the order of a finite sequence map
-reverseSeqMap :: Integer     -- ^ Size of the sequence map
-              -> SeqMap sym a
-              -> SeqMap sym a
+reverseSeqMap :: Backend sym =>
+  Integer {- ^ Size of the sequence map -} ->
+  SeqMap sym a ->
+  SeqMap sym a
 reverseSeqMap n vals = IndexSeqMap $ \i -> lookupSeqMap vals (n - 1 - i)
 
 updateSeqMap :: SeqMap sym a -> Integer -> SEval sym a -> SeqMap sym a
 updateSeqMap (UpdateSeqMap m sm) i x = UpdateSeqMap (Map.insert i x m) sm
-updateSeqMap (IndexSeqMap f) i x = UpdateSeqMap (Map.singleton i x) f
+updateSeqMap xs i x = UpdateSeqMap (Map.singleton i x) xs
 
 -- | Concatenate the first @n@ values of the first sequence map onto the
 --   beginning of the second sequence map.
-concatSeqMap :: Integer -> SeqMap sym a -> SeqMap sym a -> SeqMap sym a
+concatSeqMap :: Backend sym => Integer -> SeqMap sym a -> SeqMap sym a -> SeqMap sym a
 concatSeqMap n x y =
     IndexSeqMap $ \i ->
        if i < n
@@ -128,14 +148,14 @@ concatSeqMap n x y =
 -- | Given a number @n@ and a sequence map, return two new sequence maps:
 --   the first containing the values from @[0..n-1]@ and the next containing
 --   the values from @n@ onward.
-splitSeqMap :: Integer -> SeqMap sym a -> (SeqMap sym a, SeqMap sym a)
+splitSeqMap :: Backend sym => Integer -> SeqMap sym a -> (SeqMap sym a, SeqMap sym a)
 splitSeqMap n xs = (hd,tl)
   where
   hd = xs
   tl = IndexSeqMap $ \i -> lookupSeqMap xs (i+n)
 
 -- | Drop the first @n@ elements of the given 'SeqMap'.
-dropSeqMap :: Integer -> SeqMap sym a -> SeqMap sym a
+dropSeqMap :: Backend sym => Integer -> SeqMap sym a -> SeqMap sym a
 dropSeqMap 0 xs = xs
 dropSeqMap n xs = IndexSeqMap $ \i -> lookupSeqMap xs (i+n)
 
@@ -146,44 +166,44 @@ delaySeqMap sym xs =
 
 -- | Given a sequence map, return a new sequence map that is memoized using
 --   a finite map memo table.
-memoMap :: Backend sym => sym -> SeqMap sym a -> SEval sym (SeqMap sym a)
-memoMap sym x = do
+memoMap :: Backend sym => sym -> Nat' -> SeqMap sym a -> SEval sym (SeqMap sym a)
+
+-- Sequence is alreay memoized, just return it
+memoMap _sym _sz x@(MemoSeqMap{}) = pure x
+
+memoMap sym sz x = do
   stk <- sGetCallStack sym
   cache <- liftIO $ newIORef $ Map.empty
-  return $ IndexSeqMap (memo cache stk)
+  evalRef <- liftIO $ newIORef $ eval stk
+  return (MemoSeqMap sz cache evalRef)
 
   where
-  memo cache stk i = do
-    mz <- liftIO (Map.lookup i <$> readIORef cache)
-    case mz of
-      Just z  -> return z
-      Nothing -> sWithCallStack sym stk (doEval cache i)
+    eval stk i = sWithCallStack sym stk (lookupSeqMap x i)
 
-  doEval cache i = do
-    v <- lookupSeqMap x i
-    liftIO $ atomicModifyIORef' cache (\m -> (Map.insert i v m, ()))
-    return v
 
 -- | Apply the given evaluation function pointwise to the two given
 --   sequence maps.
 zipSeqMap ::
   Backend sym =>
   sym ->
   (a -> a -> SEval sym a) ->
+  Nat' ->
   SeqMap sym a ->
   SeqMap sym a ->
   SEval sym (SeqMap sym a)
-zipSeqMap sym f x y =
-  memoMap sym (IndexSeqMap $ \i -> join (f <$> lookupSeqMap x i <*> lookupSeqMap y i))
+zipSeqMap sym f sz x y =
+  memoMap sym sz (IndexSeqMap $ \i -> join (f <$> lookupSeqMap x i <*> lookupSeqMap y i))
 
 -- | Apply the given function to each value in the given sequence map
 mapSeqMap ::
   Backend sym =>
   sym ->
   (a -> SEval sym a) ->
-  SeqMap sym a -> SEval sym (SeqMap sym a)
-mapSeqMap sym f x =
-  memoMap sym (IndexSeqMap $ \i -> f =<< lookupSeqMap x i)
+  Nat' ->
+  SeqMap sym a ->
+  SEval sym (SeqMap sym a)
+mapSeqMap sym f sz x =
+  memoMap sym sz (IndexSeqMap $ \i -> f =<< lookupSeqMap x i)
 
 
 {-# INLINE mergeSeqMap #-}
@@ -213,7 +233,7 @@ shiftSeqByInteger sym merge reindex zro m xs idx
   | Just j <- integerAsLit sym idx = shiftOp xs j
   | otherwise =
       do (n, idx_bits) <- enumerateIntBits sym m idx
-         barrelShifter sym merge shiftOp xs n (map BitIndexSegment idx_bits)
+         barrelShifter sym merge shiftOp m xs n (map BitIndexSegment idx_bits)
  where
    shiftOp vs shft =
      pure $ indexSeqMap $ \i ->
@@ -231,6 +251,7 @@ data IndexSegment sym
   Concrete ->
   (SBit Concrete -> a -> a -> SEval Concrete a) ->
   (SeqMap Concrete a -> Integer -> SEval Concrete (SeqMap Concrete a)) ->
+  Nat' ->
   SeqMap Concrete a ->
   Integer ->
   [IndexSegment Concrete] ->
@@ -241,11 +262,12 @@ barrelShifter :: Backend sym =>
   (SBit sym -> a -> a -> SEval sym a) ->
   (SeqMap sym a -> Integer -> SEval sym (SeqMap sym a))
      {- ^ concrete shifting operation -} ->
+  Nat' {- ^ Size of the map being shifted -} ->
   SeqMap sym a {- ^ initial value -} ->
   Integer {- Number of bits in shift amount -} ->
   [IndexSegment sym]  {- ^ segments of the shift amount, in big-endian order -} ->
   SEval sym (SeqMap sym a)
-barrelShifter sym mux shift_op x0 n0 bs0
+barrelShifter sym mux shift_op sz x0 n0 bs0
   | n0 >= toInteger (maxBound :: Int) =
       liftIO (X.throw (UnsupportedSymbolicOp ("Barrel shifter with too many bits in shift amount: " ++ show n0)))
   | otherwise = go x0 (fromInteger n0) bs0
@@ -273,5 +295,5 @@ barrelShifter sym mux shift_op x0 n0 bs0
            go x_shft n' bs
       Nothing ->
         do x_shft <- shift_op x (bit n')
-           x' <- memoMap sym (mergeSeqMap sym mux b x_shft x)
+           x' <- memoMap sym sz (mergeSeqMap sym mux b x_shft x)
            go x' n' bs

src/Cryptol/Backend/WordValue.hs

@@ -66,7 +66,7 @@ import Cryptol.Backend.Concrete (Concrete(..))
 import Cryptol.Backend.Monad (EvalError(..))
 import Cryptol.Backend.SeqMap
 
-import Cryptol.TypeCheck.Solver.InfNat(widthInteger)
+import Cryptol.TypeCheck.Solver.InfNat(widthInteger, Nat'(..))
 
 -- | Force the evaluation of a word value
 forceWordValue :: Backend sym => WordValue sym -> SEval sym ()
@@ -258,19 +258,19 @@ wordValLogicOp _sym _ wop (WordVal w1) (WordVal w2) = WordVal <$> wop w1 w2
 wordValLogicOp sym bop wop (WordVal w1) (BitmapVal n2 packed2 bs2) =
   isReady sym packed2 >>= \case
     Just w2 -> WordVal <$> wop w1 w2
-    Nothing -> bitmapWordVal sym n2 =<< zipSeqMap sym bop (unpackBitmap sym w1) bs2
+    Nothing -> bitmapWordVal sym n2 =<< zipSeqMap sym bop (Nat n2) (unpackBitmap sym w1) bs2
 
 wordValLogicOp sym bop wop (BitmapVal n1 packed1 bs1) (WordVal w2) =
   isReady sym packed1 >>= \case
     Just w1 -> WordVal <$> wop w1 w2
-    Nothing -> bitmapWordVal sym n1 =<< zipSeqMap sym bop bs1 (unpackBitmap sym w2)
+    Nothing -> bitmapWordVal sym n1 =<< zipSeqMap sym bop (Nat n1) bs1 (unpackBitmap sym w2)
 
 wordValLogicOp sym bop wop (BitmapVal n1 packed1 bs1) (BitmapVal _n2 packed2 bs2) =
   do r1 <- isReady sym packed1
      r2 <- isReady sym packed2
      case (r1,r2) of
        (Just w1, Just w2) -> WordVal <$> wop w1 w2
-       _ -> bitmapWordVal sym n1 =<< zipSeqMap sym bop bs1 bs2
+       _ -> bitmapWordVal sym n1 =<< zipSeqMap sym bop (Nat n1) bs1 bs2
 
 wordValLogicOp sym bop wop (ThunkWordVal _ m1) w2 =
   do w1 <- m1
@@ -293,7 +293,7 @@ wordValUnaryOp sym bop wop (ThunkWordVal _ m) = wordValUnaryOp sym bop wop =<< m
 wordValUnaryOp sym bop wop (BitmapVal n packed xs) =
   isReady sym packed >>= \case
     Just w  -> WordVal <$> wop w
-    Nothing -> bitmapWordVal sym n =<< mapSeqMap sym bop xs
+    Nothing -> bitmapWordVal sym n =<< mapSeqMap sym bop (Nat n) xs
 
 {-# SPECIALIZE joinWords ::
   Concrete ->
@@ -537,7 +537,7 @@ shiftWordByInteger sym wop reindex x idx =
             Nothing ->
               do (numbits, idx_bits) <- enumerateIntBits' sym n idx
                  bitmapWordVal sym n =<<
-                   barrelShifter sym (iteBit sym) shiftOp bs0 numbits (map BitIndexSegment idx_bits)
+                   barrelShifter sym (iteBit sym) shiftOp (Nat n) bs0 numbits (map BitIndexSegment idx_bits)
 
  where
    shiftOp vs shft =
@@ -577,7 +577,7 @@ shiftWordByWord sym wop reindex x idx =
               bitmapWordVal sym n =<< shiftOp bs0 j
             _ ->
               do idx_segs <- enumerateIndexSegments sym idx
-                 bitmapWordVal sym n =<< barrelShifter sym (iteBit sym) shiftOp bs0 n idx_segs
+                 bitmapWordVal sym n =<< barrelShifter sym (iteBit sym) shiftOp (Nat n) bs0 n idx_segs
 
  where
    shiftOp vs shft =
@@ -654,12 +654,13 @@ shiftSeqByWord  ::
   (SBit sym -> a -> a -> SEval sym a) ->
   (Integer -> Integer -> Maybe Integer) ->
   SEval sym a ->
+  Nat' ->
   SeqMap sym a ->
   WordValue sym ->
   SEval sym (SeqMap sym a)
-shiftSeqByWord sym merge reindex zro xs idx =
+shiftSeqByWord sym merge reindex zro sz xs idx =
   do idx_segs <- enumerateIndexSegments sym idx
-     barrelShifter sym merge shiftOp xs (wordValueSize sym idx) idx_segs
+     barrelShifter sym merge shiftOp sz xs (wordValueSize sym idx) idx_segs
  where
    shiftOp vs shft =
      pure $ indexSeqMap $ \i ->
@@ -754,7 +755,7 @@ mergeBitmaps :: Backend sym =>
   SeqMap sym (SBit sym) ->
   SEval sym (WordValue sym)
 mergeBitmaps sym c sz bs1 bs2 =
-  do bs <- memoMap sym (mergeSeqMap sym (iteBit sym) c bs1 bs2)
+  do bs <- memoMap sym (Nat sz) (mergeSeqMap sym (iteBit sym) c bs1 bs2)
      bitmapWordVal sym sz bs
 
 {-# INLINE mergeWord' #-}

src/Cryptol/Eval.hs

@@ -603,7 +603,7 @@ evalComp ::
   SEval sym (GenValue sym)
 evalComp sym env len elty body ms =
        do lenv <- mconcat <$> mapM (branchEnvs sym (toListEnv env)) ms
-          mkSeq sym len elty =<< memoMap sym (indexSeqMap $ \i -> do
+          mkSeq sym len elty =<< memoMap sym len (indexSeqMap $ \i -> do
               evalExpr sym (evalListEnv lenv i) body)
 
 {-# SPECIALIZE branchEnvs ::
@@ -621,32 +621,32 @@ branchEnvs ::
   ListEnv sym ->
   [Match] ->
   SEval sym (ListEnv sym)
-branchEnvs sym env matches = foldM (evalMatch sym) env matches
+branchEnvs sym env matches = snd <$> foldM (evalMatch sym) (1, env) matches
 
 {-# SPECIALIZE evalMatch ::
   (?range :: Range, ConcPrims) =>
   Concrete ->
-  ListEnv Concrete ->
+  (Integer, ListEnv Concrete) ->
   Match ->
-  SEval Concrete (ListEnv Concrete)
+  SEval Concrete (Integer, ListEnv Concrete)
   #-}
 
 -- | Turn a match into the list of environments it represents.
 evalMatch ::
   (?range :: Range, EvalPrims sym) =>
   sym ->
-  ListEnv sym ->
+  (Integer, ListEnv sym) ->
   Match ->
-  SEval sym (ListEnv sym)
-evalMatch sym lenv m = case m of
+  SEval sym (Integer, ListEnv sym)
+evalMatch sym (lsz, lenv) m = seq lsz $ case m of
 
   -- many envs
   From n l _ty expr ->
     case len of
       -- Select from a sequence of finite length.  This causes us to 'stutter'
       -- through our previous choices `nLen` times.
       Nat nLen -> do
-        vss <- memoMap sym $ indexSeqMap $ \i -> evalExpr sym (evalListEnv lenv i) expr
+        vss <- memoMap sym (Nat lsz) $ indexSeqMap $ \i -> evalExpr sym (evalListEnv lenv i) expr
         let stutter xs = \i -> xs (i `div` nLen)
         let lenv' = lenv { leVars = fmap stutter (leVars lenv) }
         let vs i = do let (q, r) = i `divMod` nLen
@@ -655,7 +655,7 @@ evalMatch sym lenv m = case m of
                         VSeq _ xs'  -> lookupSeqMap xs' r
                         VStream xs' -> lookupSeqMap xs' r
                         _           -> evalPanic "evalMatch" ["Not a list value"]
-        return $ bindVarList n vs lenv'
+        return (lsz * nLen, bindVarList n vs lenv')
 
       -- Select from a sequence of infinite length.  Note that this means we
       -- will never need to backtrack into previous branches.  Thus, we can convert
@@ -673,15 +673,17 @@ evalMatch sym lenv m = case m of
                      VSeq _ xs'  -> lookupSeqMap xs' i
                      VStream xs' -> lookupSeqMap xs' i
                      _           -> evalPanic "evalMatch" ["Not a list value"]
-        return $ bindVarList n vs lenv'
+        -- Selecting from an infinite list effectively resets the length of the
+        -- list environment, so return 1 as the length
+        return (1, bindVarList n vs lenv')
 
     where
       len  = evalNumType (leTypes lenv) l
 
   -- XXX we don't currently evaluate these as though they could be recursive, as
   -- they are typechecked that way; the read environment to evalExpr is the same
   -- as the environment to bind a new name in.
-  Let d -> return $ bindVarList (dName d) (\i -> f (evalListEnv lenv i)) lenv
+  Let d -> return (lsz, bindVarList (dName d) (\i -> f (evalListEnv lenv i)) lenv)
     where
       f env =
           case dDefinition d of