prototype: add property test for union

prototypes/ScheduledMerges.hs

@@ -28,6 +28,7 @@ module ScheduledMerges (
     new,
     LookupResult (..),
     lookup, lookups,
+    Op,
     Update (..),
     update, updates,
     insert, inserts,
@@ -42,7 +43,9 @@ module ScheduledMerges (
     supplyUnionCredits,
 
     -- * Test and trace
+    MTree (..),
     logicalValue,
+    Representation,
     dumpRepresentation,
     representationShape,
     Event,
@@ -66,6 +69,7 @@ import           Control.Monad.ST
 import           Control.Tracer (Tracer, contramap, traceWith)
 import           GHC.Stack (HasCallStack, callStack)
 
+import qualified Test.QuickCheck as QC
 
 data LSM s  = LSMHandle !(STRef s Counter)
                         !(STRef s (LSMContent s))
@@ -1170,11 +1174,11 @@ expectCompletedMergingTree (MergingTree ref) = do
 -- Measurements
 --
 
-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
@@ -1201,7 +1205,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
@@ -1227,20 +1231,28 @@ 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])]
+type Representation =
+    ( Run
+    , [(Maybe (MergePolicy, MergeType, MergingRunState), [Run])]
+    , Maybe (MTree Run)
+    )
+
+dumpRepresentation :: LSM s -> ST s Representation
 dumpRepresentation (LSMHandle _ lsmr) = do
-    LSMContent wb ls _ <- readSTRef lsmr
-    ((Nothing, [wb]) :) <$> mapM dumpLevel ls
+    LSMContent wb ls ul <- readSTRef lsmr
+    levels <- mapM dumpLevel ls
+    tree <- case ul of
+      NoUnion   -> return Nothing
+      Union t _ -> Just <$> flattenTree t
+    return (wb, levels, tree)
 
 dumpLevel :: Level s -> ST s (Maybe (MergePolicy, MergeType, MergingRunState), [Run])
 dumpLevel (Level (Single r) rs) =
@@ -1252,14 +1264,17 @@ 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 :: Representation
+                    -> (Int, [([Int], [Int])], Maybe (MTree Int))
+representationShape (wb, levels, tree) =
+    (summaryRun wb, map summaryLevel levels, fmap (fmap summaryRun) tree)
+  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])
@@ -1296,3 +1311,26 @@ data EventDetail =
          mergeSize   :: Int
        }
   deriving stock Show
+
+-------------------------------------------------------------------------------
+-- 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

@@ -10,13 +10,17 @@ import           Data.STRef
 
 import           ScheduledMerges as LSM
 
+import qualified Test.QuickCheck as QC
+import           Test.QuickCheck (Property)
 import           Test.Tasty
 import           Test.Tasty.HUnit (HasCallStack, testCase)
+import           Test.Tasty.QuickCheck (testProperty, (=/=), (===))
 
 tests :: TestTree
-tests = testGroup "Unit tests"
-    [ testCase "regression_empty_run" test_regression_empty_run
-    , testCase "merge_again_with_incoming" test_merge_again_with_incoming
+tests = testGroup "Unit and property tests"
+    [ testCase "test_regression_empty_run" test_regression_empty_run
+    , testCase "test_merge_again_with_incoming" test_merge_again_with_incoming
+    , testProperty "prop_union" prop_union
     ]
 
 -- | Results in an empty run on level 2.
@@ -157,6 +161,39 @@ test_merge_again_with_incoming =
           , ([16,16,16,20,80], [])
           ]
 
+-------------------------------------------------------------------------------
+-- properties
+--
+
+-- | Supplying enough credits for the remaining debt completes the union merge.
+prop_union :: [[(LSM.Key, LSM.Op)]] -> Property
+prop_union kopss = length (filter (not . null) kopss) > 1 QC.==>
+    QC.ioProperty $ runWithTracer $ \tr ->
+      stToIO $ do
+        ts <- traverse (mkTable tr) kopss
+        t <- LSM.unions ts
+
+        debt <- LSM.remainingUnionDebt t
+        _ <- LSM.supplyUnionCredits t debt
+        debt' <- LSM.remainingUnionDebt t
+
+        rep <- dumpRepresentation t
+        return $ QC.counterexample (show (debt, debt')) $ QC.conjoin
+          [ debt =/= 0
+          , debt' === 0
+          , hasUnionWith isCompleted rep
+          ]
+  where
+    isCompleted = \case
+        MLeaf{} -> True
+        MNode{} -> False
+
+mkTable :: Tracer (ST s) Event -> [(LSM.Key, LSM.Op)] -> ST s (LSM s)
+mkTable tr ks = do
+    t <- LSM.new
+    LSM.updates tr t ks
+    return t
+
 -------------------------------------------------------------------------------
 -- tracing and expectations on LSM shape
 --
@@ -180,10 +217,19 @@ 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
         ]
+
+hasUnionWith :: (MTree Int -> Bool) -> Representation -> Property
+hasUnionWith p rep = do
+    let (_, _, shape) = representationShape rep
+    QC.counterexample "expected suitable Union" $
+      QC.counterexample (show shape) $
+        case shape of
+          Nothing -> False
+          Just t  -> p t

prototypes/ScheduledMergesTestQLS.hs

@@ -1,7 +1,5 @@
 {-# LANGUAGE TypeFamilies #-}
 
-{-# OPTIONS_GHC -Wno-orphans #-}
-
 module ScheduledMergesTestQLS (tests) where
 
 import           Control.Monad.ST
@@ -393,19 +391,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