module BTree.Merge ( mergeTrees
, mergeLeaves
, sizedProducerForTree
) where
import Control.Applicative
import Data.Foldable
import Control.Monad.State hiding (forM_)
import Control.Monad.Catch
import Data.Binary
import Control.Lens
import Pipes
import Pipes.Interleave
import Prelude hiding (sum)
import BTree.Types
import BTree.Builder
import BTree.Walk
-- | Merge trees' leaves taking ordered leaves from a set of producers.
--
-- Each producer must be annotated with the number of leaves it is
-- expected to produce. The size of the resulting tree will be at most
-- the sum of these sizes.
mergeLeaves :: (MonadMask m, MonadIO m, Functor m, Binary k, Binary e, Ord k)
=> (e -> e -> m e) -- ^ merge operation on elements
-> Order -- ^ order of merged tree
-> FilePath -- ^ name of output file
-> [(Size, Producer (BLeaf k e) m ())] -- ^ producers of leaves to merge
-> m ()
mergeLeaves append destOrder destFile producers = do
let size = sum $ map fst producers
fromOrderedToFile destOrder size destFile $
mergeM doAppend (map snd producers)
where
doAppend (BLeaf k e) (BLeaf _ e') = BLeaf k <$> append e e'
{-# INLINE mergeLeaves #-}
-- | Merge several 'LookupTrees'
--
-- This is a convenience function for merging several trees already on
-- disk. For a more flexible interface, see 'mergeLeaves'.
mergeTrees :: (MonadMask m, MonadIO m, Functor m, Binary k, Binary e, Ord k)
=> (e -> e -> m e) -- ^ merge operation on elements
-> Order -- ^ order of merged tree
-> FilePath -- ^ name of output file
-> [LookupTree k e] -- ^ trees to merge
-> m ()
mergeTrees append destOrder destFile trees = do
mergeLeaves append destOrder destFile
$ map sizedProducerForTree trees
{-# INLINE mergeTrees #-}
-- | Get a sized 'Producer' suitable for 'mergeLeaves' from a 'LookupTree'
sizedProducerForTree :: (Monad m, Binary k, Binary e)
=> LookupTree k e -- ^ a tree
-> (Size, Producer (BLeaf k e) m ())
-- ^ a sized 'Producer' suitable for passing
-- to 'mergeLeaves'
sizedProducerForTree lt = (lt ^. ltHeader . btSize, void $ walkLeaves lt)