prototype: add property test for union

prototypes/ScheduledMerges.hs

@@ -25,6 +25,7 @@ module ScheduledMerges (
     new,
     LookupResult (..),
     lookup, lookups,
+    Op,
     Update (..),
     update, updates,
     insert, inserts,
@@ -39,6 +40,7 @@ module ScheduledMerges (
     supplyUnionCredits,
 
     -- * Test and trace
+    MTree (..),
     logicalValue,
     dumpRepresentation,
     representationShape,
@@ -62,6 +64,8 @@ import           Control.Monad.ST
 import           Control.Tracer (Tracer, contramap, nullTracer, traceWith)
 import           GHC.Stack (HasCallStack, callStack)
 
+import qualified Test.QuickCheck as QC
+
 
 data LSM s  = LSMHandle !(STRef s Counter)
                         !(STRef s (LSMContent s))
@@ -1110,11 +1114,11 @@ expectCompletedMergingTree (MergingTree ref) = do
 --
 
 -- TODO: Is this useful or should we directly flatten to a Map?
-data MTree = MLeaf Run
-           | MNode MergeType [MTree]
-  deriving stock Show
+data MTree r = MLeaf r
+             | MNode MergeType [MTree r]
+  deriving stock (Eq, Foldable, Functor, Show)
 
-allLayers :: LSM s -> ST s (Buffer, [[Run]], Maybe MTree)
+allLayers :: LSM s -> ST s (Buffer, [[Run]], Maybe (MTree Run))
 allLayers (LSMHandle _ lsmr) = do
     LSMContent wb ls ul <- readSTRef lsmr
     rs <- flattenLevels ls
@@ -1141,7 +1145,7 @@ flattenMergingRun (MergingRun _ ref) = do
       CompletedMerge r    -> return [r]
       OngoingMerge _ rs _ -> return rs
 
-flattenTree :: MergingTree s -> ST s MTree
+flattenTree :: MergingTree s -> ST s (MTree Run)
 flattenTree (MergingTree ref) = do
     mts <- readSTRef ref
     case mts of
@@ -1167,20 +1171,27 @@ logicalValue lsm = do
     mergeRuns :: MergeType -> [Run] -> Run
     mergeRuns = mergek
 
-    mergeTree :: MTree -> Run
+    mergeTree :: MTree Run -> Run
     mergeTree (MLeaf r)     = r
     mergeTree (MNode mt ts) = mergeRuns mt (map mergeTree ts)
 
     justInsert (Insert v b) = Just (v, b)
     justInsert  Delete      = Nothing
     justInsert (Mupsert v)  = Just (v, Nothing)
 
--- TODO: Consider MergingTree, or just remove this function? It's unused.
 dumpRepresentation :: LSM s
-                   -> ST s [(Maybe (MergePolicy, MergeType, MergingRunState), [Run])]
+                   -> ST s
+                        ( Run
+                        , [(Maybe (MergePolicy, MergeType, MergingRunState), [Run])]
+                        , Maybe (MTree Run)
+                        )
 dumpRepresentation (LSMHandle _ lsmr) = do
-    LSMContent wb ls _ <- readSTRef lsmr
-    ((Nothing, [wb]) :) <$> mapM dumpLevel ls
+    LSMContent wb ls ul <- readSTRef lsmr
+    levels <- mapM dumpLevel ls
+    mtree <- case ul of
+      NoUnion   -> return Nothing
+      Union t _ -> Just <$> flattenTree t
+    return (wb, levels, mtree)
 
 dumpLevel :: Level s -> ST s (Maybe (MergePolicy, MergeType, MergingRunState), [Run])
 dumpLevel (Level (Single r) rs) =
@@ -1192,15 +1203,41 @@ dumpLevel (Level (Merging mp (MergingRun mt ref)) rs) = do
 -- For each level:
 -- 1. the runs involved in an ongoing merge
 -- 2. the other runs (including completed merge)
-representationShape :: [(Maybe (MergePolicy, MergeType, MergingRunState), [Run])]
-                    -> [([Int], [Int])]
-representationShape =
-    map $ \(mmr, rs) ->
+representationShape :: (Run, [(Maybe (MergePolicy, MergeType, MergingRunState), [Run])], Maybe (MTree Run))
+                    -> (Int, [([Int], [Int])], Maybe (MTree Int))
+representationShape (wb, levels, mtree) =
+    (summaryRun wb, map summaryLevel levels, fmap (fmap summaryRun) mtree)
+  where
+    summaryLevel (mmr, rs) =
       let (ongoing, complete) = summaryMR mmr
       in (ongoing, complete <> map summaryRun rs)
-  where
+
     summaryRun = runSize
+
     summaryMR = \case
       Nothing                          -> ([], [])
       Just (_, _, CompletedMerge r)    -> ([], [summaryRun r])
       Just (_, _, OngoingMerge _ rs _) -> (map summaryRun rs, [])
+
+-------------------------------------------------------------------------------
+-- Arbitrary
+--
+
+instance QC.Arbitrary Key where
+  arbitrary = K <$> QC.arbitrarySizedNatural
+  shrink (K v) = K <$> QC.shrink v
+
+instance QC.Arbitrary Value where
+  arbitrary = V <$> QC.arbitrarySizedNatural
+  shrink (V v) = V <$> QC.shrink v
+
+instance QC.Arbitrary Blob where
+  arbitrary = B <$> QC.arbitrarySizedNatural
+  shrink (B v) = B <$> QC.shrink v
+
+instance (QC.Arbitrary v, QC.Arbitrary b) => QC.Arbitrary (Update v b) where
+  arbitrary = QC.frequency
+      [ (3, Insert <$> QC.arbitrary <*> QC.arbitrary)
+      , (1, Mupsert <$> QC.arbitrary)
+      , (1, pure Delete)
+      ]

prototypes/ScheduledMergesTest.hs

@@ -1,7 +1,7 @@
 module ScheduledMergesTest (tests) where
 
 import           Control.Exception
-import           Control.Monad (replicateM_, when)
+import           Control.Monad (forM, replicateM_, when)
 import           Control.Monad.ST
 import           Control.Tracer (Tracer (Tracer))
 import qualified Control.Tracer as Tracer
@@ -12,11 +12,15 @@ import           ScheduledMerges as LSM
 
 import           Test.Tasty
 import           Test.Tasty.HUnit (HasCallStack, testCase)
+import           Test.Tasty.QuickCheck (testProperty)
+import           Test.QuickCheck (Property)
+import qualified Test.QuickCheck as QC
 
 tests :: TestTree
 tests = testGroup "Unit tests"
     [ testCase "regression_empty_run" test_regression_empty_run
     , testCase "merge_again_with_incoming" test_merge_again_with_incoming
+    , testProperty "union" prop_union
     ]
 
 -- | Results in an empty run on level 2.
@@ -140,6 +144,34 @@ test_merge_again_with_incoming =
           , ([16,16,16,20,80], [])
           ]
 
+-------------------------------------------------------------------------------
+-- properties
+--
+
+-- | Supplying enough credits for the remaining debt fully merges the tree.
+prop_union :: [[(LSM.Key, LSM.Op)]] -> Property
+prop_union kopss = length (filter (not . null) kopss) > 1 QC.==>
+    QC.ioProperty $ runWithTracer $ \tr ->
+      stToIO $ do
+        ts <- forM kopss $ \kops -> do
+          t <- LSM.new
+          LSM.updates tr t kops
+          return t
+
+        t <- LSM.unions ts
+
+        expectUnionWith t (not . isCompleted)
+
+        debt <- LSM.remainingUnionDebt t
+        _ <- LSM.supplyUnionCredits t debt
+
+        -- assert that the union merge now has been completed
+        expectUnionWith t isCompleted
+  where
+    isCompleted = \case
+        MLeaf{} -> True
+        MNode{} -> False
+
 -------------------------------------------------------------------------------
 -- tracing and expectations on LSM shape
 --
@@ -163,10 +195,17 @@ instance Exception TracedException where
 
 expectShape :: HasCallStack => LSM s -> Int -> [([Int], [Int])] -> ST s ()
 expectShape lsm expectedWb expectedLevels = do
-    let expected = ([], [expectedWb]) : expectedLevels
+    let expected = (expectedWb, expectedLevels, Nothing)
     shape <- representationShape <$> dumpRepresentation lsm
     when (shape /= expected) $
       error $ unlines
         [ "expected shape: " <> show expected
         , "actual shape:   " <> show shape
         ]
+
+expectUnionWith :: HasCallStack => LSM s -> (MTree Int -> Bool) -> ST s ()
+expectUnionWith lsm p = do
+    (_, _, shape) <- representationShape <$> dumpRepresentation lsm
+    case shape of
+      Nothing -> error "expected Union, found NoUnion"
+      Just t  -> when (not (p t)) $ error $ "expectation on Union failed: " ++ show t

prototypes/ScheduledMergesTestQLS.hs

@@ -393,19 +393,3 @@ runModel action ctx m =
 
     lookUpKeyVar :: ModelVar Model Key -> Key
     lookUpKeyVar var = case lookupVar ctx var of MInsert k -> k
-
--------------------------------------------------------------------------------
--- Instances
---
-
-instance Arbitrary Key where
-  arbitrary = K <$> arbitrarySizedNatural
-  shrink (K v) = K <$> shrink v
-
-instance Arbitrary Value where
-  arbitrary = V <$> arbitrarySizedNatural
-  shrink (V v) = V <$> shrink v
-
-instance Arbitrary Blob where
-  arbitrary = B <$> arbitrarySizedNatural
-  shrink (B v) = B <$> shrink v