content: 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)
  