added NewUTxO

libs/compact-map/compact-map.cabal

@@ -33,7 +33,8 @@ library
                      , Data.Compact.HashMap
                      , Data.Compact.VMap
                      , Data.Compact.SmallArray
-  other-modules:       Data.Compact.KVVector
+                     , Data.Compact.SplitMap
+  other-modules:     , Data.Compact.KVVector
   build-depends:       base >=4.11 && <5
                      , cardano-binary
                      , cardano-prelude

libs/compact-map/src/Data/Compact/KeyMap.hs

@@ -525,9 +525,49 @@ unionWith comb x y = union4 0 (\_k a b -> comb a b) x y
 union :: KeyMap v -> KeyMap v -> KeyMap v
 union x y = union4 0 (\_k a _b -> a) x y
 
+
 -- ===========================================
 -- intersection operators
 
+-- | Get the largest key, NOT the largest value
+getMax :: KeyMap v -> Maybe (Key, v)
+getMax Empty = Nothing
+getMax (Leaf k v) = Just (k, v)
+getMax (One _ x) = getMax x
+getMax (Two _ _ y) = getMax y
+getMax (BitmapIndexed _ arr) = getMax (index arr (isize arr - 1))
+getMax (Full arr) = getMax (index arr (isize arr - 1))
+
+minViewWithKey :: KeyMap v -> Maybe ((Key, v), KeyMap v)
+minViewWithKey _x = Nothing
+
+-- ==================================================
+
+-- | The (key,value) pairs (i.e. a subset) of 'h1' where key is in the domain of both 'h1' and 'h2'
+intersect :: KeyMap v -> KeyMap v -> KeyMap v
+intersect map1 map2 =
+  case maxMinOf map1 map2 of
+    Nothing -> Empty
+    Just k -> leapfrog k map1 map2 Empty
+
+-- | Accumulate a new Key map, by adding the key value pairs to 'ans', for
+--   the Keys that appear in both maps 'x' and 'y'. The key 'k' should
+--   be the smallest key in either 'x' or 'y', used to get started.
+leapfrog :: Key -> KeyMap v -> KeyMap v -> KeyMap v -> KeyMap v
+leapfrog k x y ans =
+  case (lub k x, lub k y) of
+    (Just (k1, v1, h1), Just (k2, _, h2)) ->
+      case maxMinOf h1 h2 of
+        Just k3 -> leapfrog k3 h1 h2 (if k1 == k2 then insert k1 v1 ans else ans)
+        Nothing -> (if k1 == k2 then insert k1 v1 ans else ans)
+    _ -> ans
+
+-- | Get the larger of the two min keys of 'x' and 'y'. Nothing if either is Empty.
+maxMinOf :: KeyMap v1 -> KeyMap v2 -> Maybe Key
+maxMinOf x y = case (getMin x, getMin y) of
+  (Just (k1, _), Just (k2, _)) -> Just (max k1 k2)
+  _ -> Nothing
+
 intersect3 :: Int -> (Key -> u -> v -> w) -> KeyMap u -> KeyMap v -> KeyMap w
 intersect3 _ _ Empty Empty = Empty
 intersect3 n combine x y = case3 Empty leafF1 arrayF1 x

libs/compact-map/src/Data/Compact/SplitMap.hs

@@ -0,0 +1,148 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Data.Compact.SplitMap where
+
+import Data.Compact.KeyMap (Key, KeyMap)
+import qualified Data.Compact.KeyMap as KeyMap
+import qualified Data.IntMap as IntMap
+import Data.IntMap.Strict (IntMap)
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Prelude hiding (lookup)
+
+class Split k where
+  split :: k -> (Int, Key)
+  join :: Int -> Key -> k
+
+data SplitMap k v where
+  SplitMap :: Split k => IntMap (KeyMap v) -> SplitMap k v
+
+empty :: forall k v. Split k => SplitMap k v
+empty = SplitMap IntMap.empty
+
+insertWithKey :: forall k v. (k -> v -> v -> v) -> k -> v -> SplitMap k v -> SplitMap k v
+insertWithKey combine k v (SplitMap imap) = SplitMap (IntMap.insertWith combine2 n (KeyMap.insert key v KeyMap.Empty) imap)
+  where
+    (n, key) = split k
+    combine2 :: KeyMap v -> KeyMap v -> KeyMap v
+    combine2 km1 km2 = KeyMap.unionWith (combine k) km1 km2
+
+insertWith :: forall k v. (v -> v -> v) -> k -> v -> SplitMap k v -> SplitMap k v
+insertWith comb k v mp = insertWithKey (\_ x y -> comb x y) k v mp
+
+insert :: forall k v. k -> v -> SplitMap k v -> SplitMap k v
+insert k v mp = insertWithKey (\_k v1 _v2 -> v1) k v mp
+
+delete :: k -> SplitMap k v -> SplitMap k v
+delete k (SplitMap imap) = SplitMap (IntMap.update fix n imap)
+  where
+    (n, key) = split k
+    fix keymap = case KeyMap.delete key keymap of
+      KeyMap.Empty -> Nothing
+      !other -> Just other
+
+lookup :: k -> SplitMap k v -> Maybe v
+lookup k (SplitMap imap) =
+  case IntMap.lookup n imap of
+    Nothing -> Nothing
+    Just keymap -> KeyMap.lookupHM key keymap
+  where
+    (n, key) = split k
+
+mapWithKey :: forall k v u. (k -> v -> u) -> SplitMap k v -> SplitMap k u
+mapWithKey f (SplitMap imap) = SplitMap (IntMap.mapWithKey g imap)
+  where
+    g :: Int -> KeyMap v -> KeyMap u
+    g n kmap = KeyMap.mapWithKey (\key v -> f (join n key) v) kmap
+
+unionWithKey :: forall k v. (k -> v -> v -> v) -> SplitMap k v -> SplitMap k v -> SplitMap k v
+unionWithKey combine (SplitMap imap1) (SplitMap imap2) = SplitMap (IntMap.unionWithKey comb imap1 imap2)
+  where
+    comb :: Int -> KeyMap v -> KeyMap v -> KeyMap v
+    comb n x y = KeyMap.unionWithKey (\key v1 v2 -> combine (join n key) v1 v2) x y
+
+unionWith :: forall k v. (v -> v -> v) -> SplitMap k v -> SplitMap k v -> SplitMap k v
+unionWith combine (SplitMap imap1) (SplitMap imap2) = SplitMap (IntMap.unionWith comb imap1 imap2)
+  where
+    comb :: KeyMap v -> KeyMap v -> KeyMap v
+    comb x y = KeyMap.unionWith combine x y
+
+union :: forall k v. SplitMap k v -> SplitMap k v -> SplitMap k v
+union (SplitMap imap1) (SplitMap imap2) = SplitMap (IntMap.unionWith comb imap1 imap2)
+  where
+    comb :: KeyMap v -> KeyMap v -> KeyMap v
+    comb x y = KeyMap.unionWith (\v _ -> v) x y
+
+-- ============================================================================
+
+foldlWithKey' :: forall k v ans. (ans -> k -> v -> ans) -> ans -> SplitMap k v -> ans
+foldlWithKey' comb ans0 (SplitMap imap) = IntMap.foldlWithKey' comb2 ans0 imap
+  where
+    comb2 :: ans -> Int -> KeyMap v -> ans
+    comb2 ans1 n kmap = KeyMap.foldWithDescKey comb3 ans1 kmap
+      where
+        comb3 :: Key -> v -> ans -> ans
+        comb3 key v ans2 = comb ans2 (join n key) v
+
+foldrWithKey' :: forall k v ans. (k -> ans -> v -> ans) -> ans -> SplitMap k v -> ans
+foldrWithKey' comb ans0 (SplitMap imap) = IntMap.foldrWithKey' comb2 ans0 imap
+  where
+    comb2 :: Int -> KeyMap v -> ans -> ans
+    comb2 n kmap ans1 = KeyMap.foldWithAscKey comb3 ans1 kmap
+      where
+        comb3 :: ans -> Key -> v -> ans
+        comb3 ans2 key v = comb (join n key) ans2 v
+
+-- =================================================================================
+-- These 'restrictKeys' functions assume the structure holding the 'good' keys is small
+-- An alternate approach is to use cross-type 'intersection' operations
+
+restrictKeysSet :: forall k a. SplitMap k a -> Set k -> SplitMap k a
+restrictKeysSet splitmap@(SplitMap _) kset = Set.foldl' comb (SplitMap IntMap.empty) kset
+  where
+    comb :: SplitMap k a -> k -> SplitMap k a
+    comb smap k = case lookup k splitmap of
+      Nothing -> smap
+      Just a -> insert k a smap
+
+restrictKeysMap :: forall k a b. SplitMap k a -> Map k b -> SplitMap k a
+restrictKeysMap splitmap@(SplitMap _) kmap = Map.foldlWithKey' comb (SplitMap IntMap.empty) kmap
+  where
+    comb :: SplitMap k a -> k -> b -> SplitMap k a
+    comb smap k _ = case lookup k splitmap of
+      Nothing -> smap
+      Just a -> insert k a smap
+
+restrictKeysSplit :: forall k a b. SplitMap k a -> SplitMap k b -> SplitMap k a
+restrictKeysSplit splitmap@(SplitMap _) ksplit = foldlWithKey' comb (SplitMap IntMap.empty) ksplit
+  where
+    comb :: SplitMap k a -> k -> b -> SplitMap k a
+    comb smap k _ = case lookup k splitmap of
+      Nothing -> smap
+      Just a -> insert k a smap
+
+-- =================================================================================
+-- These 'withoutKeys' functions assume the structure holding the 'bad' keys is small
+-- An alternate approach is to use cross-type 'intersection' operations
+
+withoutKeysSet :: forall k a. SplitMap k a -> Set k -> SplitMap k a
+withoutKeysSet splitmap@(SplitMap _) kset = Set.foldl' comb splitmap kset
+  where
+    comb :: SplitMap k a -> k -> SplitMap k a
+    comb smap k = delete k smap
+
+withoutKeysMap :: forall k a b. SplitMap k a -> Map k b -> SplitMap k a
+withoutKeysMap splitmap@(SplitMap _) kset = Map.foldlWithKey' comb splitmap kset
+  where
+    comb :: SplitMap k a -> k -> b -> SplitMap k a
+    comb smap k _ = delete k smap
+
+withoutKeysSplit :: forall k a b. SplitMap k a -> SplitMap k b -> SplitMap k a
+withoutKeysSplit splitmap@(SplitMap _) kset = foldlWithKey' comb splitmap kset
+  where
+    comb :: SplitMap k a -> k -> b -> SplitMap k a
+    comb smap k _ = delete k smap

libs/small-steps/small-steps.cabal

@@ -49,6 +49,7 @@ library
                      , Control.Provenance
                      , Control.Iterate.SetAlgebra
                      , Control.Iterate.Collect
+                     , Control.Iterate.SplitMap
                      , Control.SetAlgebra
   build-depends:       aeson
                      , ansi-wl-pprint

libs/small-steps/src/Control/Iterate/SetAlgebra.hs

@@ -30,6 +30,8 @@ import Control.DeepSeq (NFData (rnf))
 import Control.Iterate.Collect
 import Control.Monad (unless, void)
 import Data.Coders (cborError, invalidKey)
+import Data.Compact.SplitMap (Split (..), SplitMap (..))
+import qualified Data.Compact.SplitMap as SplitMap
 import Data.List (sortBy)
 import qualified Data.List as List
 import Data.Map.Internal (Map (..), link, link2)
@@ -595,6 +597,7 @@ data BaseRep f k v where
   ListR :: Basic List => BaseRep List k v
   SingleR :: Basic Single => BaseRep Single k v
   BiMapR :: (Basic (BiMap v), Ord v) => BaseRep (BiMap v) k v
+  SplitR :: Split k => BaseRep SplitMap k v
 
 -- ==========================================================================
 -- The most basic operation of iteration, where (Iter f) is to use the 'nxt'
@@ -933,12 +936,13 @@ lift f = Fun (Lift f) f
 -- to be applied to a collection built by iterating over a Query. This produces the keys in
 -- ascending order, with no duplicate keys. So we do not need to specify how to merge values.
 -- =============================================================================================
-materialize :: (Ord k) => BaseRep f k v -> Collect (k, v) -> f k v
+materialize :: forall f k v. (Ord k) => BaseRep f k v -> Collect (k, v) -> f k v
 materialize ListR x = fromPairs (\l _r -> l) (runCollect x [] (:))
 materialize MapR x = runCollect x Map.empty (\(k, v) ans -> Map.insert k v ans)
 materialize SetR x = Sett (runCollect x Set.empty (\(k, _) ans -> Set.insert k ans))
 materialize BiMapR x = runCollect x biMapEmpty (\(k, v) ans -> addpair k v ans)
 materialize SingleR x = runCollect x Fail (\(k, v) _ignore -> Single k v)
+materialize SplitR x = runCollect x (SplitMap.empty @k) (\(k, v) ans -> SplitMap.insert k v ans)
 
 -- ================================================================================
 -- On the flip side, a witness can be used to specifiy how to build a datatype from
@@ -954,12 +958,13 @@ addp combine (k, v) xs = addkv (k, v) xs combine
 -- later in the list override ones earlier in the list, and comb =
 -- (\ earlier later -> earlier) will keep the value that appears first in the list
 
-fromList :: Ord k => BaseRep f k v -> (v -> v -> v) -> [(k, v)] -> f k v
+fromList :: forall f k v. Ord k => BaseRep f k v -> (v -> v -> v) -> [(k, v)] -> f k v
 fromList MapR combine xs = Map.fromListWith combine xs
 fromList ListR combine xs = fromPairs combine xs
 fromList SetR combine xs = foldr (addp combine) (Sett (Set.empty)) xs
 fromList BiMapR combine xs = biMapFromList combine xs
 fromList SingleR combine xs = foldr (addp combine) Fail xs
+fromList SplitR combine xs = foldr (\(k, v) a -> SplitMap.insertWith combine k v a) (SplitMap.empty @k) xs
 
 -- =========================================================================================
 -- Now we make an iterator that collects triples, on the intersection
@@ -1115,6 +1120,7 @@ runSet e = run (compile e)
 run :: (Ord k) => (Query k v, BaseRep f k v) -> f k v
 run (BaseD SetR x, SetR) = x -- If it is already data (BaseD)
 run (BaseD MapR x, MapR) = x -- and in the right form (the BaseRep's match)
+run (BaseD SplitR x, SplitR) = x
 run (BaseD SingleR x, SingleR) = x -- just return the data
 run (BaseD BiMapR x, BiMapR) = x -- only need to materialize data
 run (BaseD ListR x, ListR) = x -- if the forms do not match.
@@ -1470,9 +1476,13 @@ guardQ ::
   Query k v
 guardQ x p = GuardD (query x) p
 
--- Don't know why this won't type check
--- diffQ :: (Ord k, HasQuery concrete1 k v, HasQuery concrete2 k u) => concrete1 -> concrete2 -> Query k v
--- diffQ = \ x y -> DiffD (query x) (query y)
+diffQ ::
+  forall k v u concrete1 concrete2.
+  (Ord k, HasQuery (concrete1 k v) k v, HasQuery (concrete2 k u) k u) =>
+  (concrete1 k v) ->
+  (concrete2 k u) ->
+  Query k v
+diffQ x y = DiffD (query x) (query @(concrete2 k u) @k @u y)
 
 class HasQuery concrete k v where
   query :: concrete -> Query k v
@@ -1505,6 +1515,7 @@ instance Show (BaseRep f k v) where
   show ListR = "List"
   show SingleR = "Single"
   show BiMapR = "BiMap"
+  show SplitR = "SplitMap"
 
 instance Show (Exp t) where
   show (Base MapR x) = "Map(" ++ show (Map.size x) ++ ")?"

libs/small-steps/src/Control/Iterate/SplitMap.hs

@@ -0,0 +1,74 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module Control.Iterate.SplitMap where
+
+import Control.Iterate.Collect
+import Control.Iterate.SetAlgebra
+import qualified Data.Compact.KeyMap as KeyMap
+import Data.Compact.SplitMap (Split (..), SplitMap (..))
+import qualified Data.Compact.SplitMap as SplitMap
+import qualified Data.IntMap as IntMap
+import qualified Data.Set as Set
+
+-- ======================================
+
+-- | Insert a KeyMap into an IntMap, unless it is empty, if so, return the IntMap unchanged
+--   we assume the Int 'n' is not already in the IntMap 'imap', and each call site should have this invariant.
+insertNormForm :: Split k => IntMap.Key -> KeyMap.KeyMap v -> IntMap.IntMap (KeyMap.KeyMap v) -> SplitMap k v
+insertNormForm _ (KeyMap.Empty) imap = SplitMap imap
+insertNormForm n kmap imap = SplitMap (IntMap.insert n kmap imap)
+
+instance Iter SplitMap where
+  nxt (SplitMap imap) =
+    case IntMap.minViewWithKey imap of
+      Nothing -> none
+      Just ((n, kmap), imap2) ->
+        case KeyMap.minViewWithKey kmap of
+          Nothing -> none -- This should never happen, every 'n' should have at least one 'key'
+          Just ((key, v), kmap2) -> one (join n key, v, insertNormForm n kmap2 imap2)
+
+  lub k (SplitMap imap) =
+    let (n, key) = split k
+     in case IntMap.splitLookup n imap of
+          (_, Just kmap, imap2) ->
+            case KeyMap.lub key kmap of
+              Nothing -> none -- This should never happen, every 'n' should have at least one 'key'
+              Just (key2, v, kmap2) -> one (join n key2, v, insertNormForm n kmap2 imap2)
+          (_, Nothing, imap3) -> nxt (SplitMap imap3)
+
+  isnull (SplitMap x) = IntMap.null x
+
+  haskey k x =
+    case SplitMap.lookup k x of
+      Nothing -> False
+      Just _ -> True
+
+  lookup k x = SplitMap.lookup k x
+
+  element k x =
+    case SplitMap.lookup k x of
+      Nothing -> none
+      Just _ -> one ()
+
+instance Basic SplitMap where
+  addpair k v x = SplitMap.insert k v x
+  addkv (k, v) x comb = SplitMap.insertWith comb k v x
+  removekey = SplitMap.delete
+  domain smap = SplitMap.foldlWithKey' accum Set.empty smap
+    where
+      accum ans k _ = Set.insert k ans
+  range smap = SplitMap.foldlWithKey' accum Set.empty smap
+    where
+      accum ans _ v = Set.insert v ans
+
+instance (Ord k, Split k) => HasExp (SplitMap k v) (SplitMap k v) where
+  toExp x = Base SplitR x
+
+instance (Ord k, Split k) => HasQuery (SplitMap k v) k v where
+  query xs = BaseD SplitR xs
+
+instance Embed (SplitMap k v) (SplitMap k v) where
+  toBase xs = xs
+  fromBase xs = xs