module Data.List.Merge where

-- | Ordinary Ord-based sorted list merge.
--
-- TODO: verify fusion.
mergeL :: Ord a => [a] -> [a] -> [a]
mergeL as bs = mergeLists (mergeData compare as bs) const as bs

-- | Merge lists based on pre-computed comparison results.  Use 'mergeData' to
-- perform the comparisons.
mergeLists :: [(Int,Ordering)] -- ^ comparison results.
              -> (a -> a -> a) -- ^ combining function applied when 'EQ' is encountered.
              -> [a] -- ^ sorted list
              -> [a] -- ^ sorted list
              -> [a] -- ^ merged sorted list
mergeLists ((n,LT):os) f xs ys = ls ++ mergeLists os f xs' ys
 where
    (ls,xs') = splitAt n xs
mergeLists ((n,EQ):os) f xs ys = es ++ mergeLists os f xs' ys'
 where
    (les,xs') = splitAt n xs
    (res,ys') = splitAt n ys
    es = zipWith f les res
mergeLists ((n,GT):os) f xs ys = gs ++ mergeLists os f xs ys'
 where
    (gs,ys') = splitAt n ys
mergeLists [] _ [] ys = ys
mergeLists [] _ xs [] = xs
mergeLists [] _ _  _  = error "mergeLists: insufficient data." -- xs ++ ys

-- | Inverse to 'mergeLists': given a list of comparison results, partition a
-- list into the parts necessary for 'mergeLists' to recreate it.
splitLists :: [(Int,Ordering)] -> [a] -> ([a],[a])
splitLists ((n,LT):os) xs = (ls ++ lls, rrs)
 where
    (ls,xs') = splitAt n xs
    (lls,rrs) = splitLists os xs'
splitLists ((n,EQ):os) xs = (es ++ lls, es ++ rrs)
 where
    (es,xs') = splitAt n xs
    (lls,rrs) = splitLists os xs'
splitLists ((n,GT):os) xs = (lls, rs ++ rrs)
 where
    (rs,xs') = splitAt n xs
    (lls,rrs) = splitLists os xs'
splitLists [] xs = (xs,xs)


-- | mergeData
--
--   > mergeData compare [1,3,5] [2,2,4,6]  ==> [(1,LT),(2,GT),(1,LT),(1,GT),(1,LT),(1,GT)]
--
--   Given a comparison function and two sorted lists, 'mergeData' will return
--   a RLE compressed (run-length encoded) list of the comparison results
--   encountered while merging the lists.
--
--   This data is enough information to perform the merge without doing the
--   comparisons or to reverse a merged list back to two sorted lists.
--
--   When one list is exhausted, the length of the remaining list is returned
--   as a run-length for LT or GT depending on whether the left list or the
--   right list has elements.
mergeData :: (a -> a -> Ordering) -> [a] -> [a] -> [(Int,Ordering)]
mergeData comp (x:xs) (y:ys)
  | comp x y == LT = case mergeData comp xs (y:ys) of
                        (n,LT):zs -> let n'=n+1 in n' `seq` (n',LT):zs
                        zs        -> (1,LT):zs
  | comp x y == EQ = case mergeData comp xs ys of
                        (n,EQ):zs -> let n'=n+1 in n' `seq` (n',EQ):zs
                        zs        -> (1,EQ):zs
  | otherwise      = case mergeData comp (x:xs) ys of
                        (n,GT):zs -> let n'=n+1 in n' `seq` (n',GT):zs
                        zs        -> (1,GT):zs
mergeData _    [] [] = []
mergeData _    [] ys = (length ys, GT) : []
mergeData _    xs [] = (length xs, LT) : []