ZipAlign.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
import Data.Foldable as F
import Data.Traversable as T
import Data.Monoid
import Control.Monad.State
import Control.Applicative
import Control.Monad.Writer
import Control.Applicative.Compose
import Data.Either
import qualified Data.Map as M

{- | filtering is not possible with other classes?

Laws might be:

>  id == fst . partition . fmap Left
>  id == snd . partition . fmap Right

>  partition == fmap swap . partition . fmap swap
> where
>  swap (Left x) = Right x
>  swap (Right x) = Left x

> partitionEithers . F.toList == \x -> case partition x of (a,b) -> (F.toList a, F.toList b)

-}
class Traversable f => Partition f where partition :: f (Either a b) -> (f a, f b)

instance Partition [] where partition = partitionEithers
instance Partition Maybe where
    partition (Just (Left x)) = (Just x, Nothing)
    partition (Just (Right x)) = (Nothing, Just x)
    partition _ = (Nothing, Nothing)
instance Partition (M.Map k) where
    partition x = let
        (l,r) = M.partition (\a -> case a of Left {} -> True; _ -> False) x
      in (fmap (\(Left x) -> x) l, fmap (\(Right x) -> x) r)


newtype M f a = M ((f :+: Maybe) a) deriving (Functor, Applicative)

deriving instance Show ((f :+: Maybe) a) => Show (M f a)
deriving instance Show (f (g a)) => Show ((f :+: g) a)

instance Foldable f => Foldable (M f) where
    foldMap f (M (Compose x)) = foldMap (maybe mempty f) x

instance Traversable f => Traversable (M f) where
    traverse f (M (Compose x)) = fmap (M . Compose) $ traverse (\y -> case y of
        Just y -> fmap Just (f y)
        Nothing -> pure Nothing) x

instance Traversable f => Partition (M f) where
    partition (M (Compose x)) =
        (M $ Compose $ fmap (\x -> case x of
            Just (Left y) -> Just y
            _ -> Nothing) x,
        M $ Compose $ fmap (\x -> case x of
                Just (Right y) -> Just y
                _ -> Nothing) x)


zipAlignG2 :: Traversable f => f a -> f a1 -> Zip f (Maybe a) (Maybe a1)
zipAlignG2 x y = unM $ zipAlignG (M (Compose (fmap Just x))) (M (Compose (fmap Just y)))

-- not really clear it's a good idea to distribute the Just over (,)
unM (R (M x)) = R $ decompose x
unM (L (M x)) = L $ decompose x
unM (C (M x)) = C $ fmap (\ x -> case x of
            Just (a,b) -> (Just a, Just b)
            _ -> (Nothing, Nothing)) $ decompose x

test1 = zipAlignG (M $ Compose $ [Just 5]) (M $ Compose $ [Nothing])

zipAlignG as bs
    | (partition -> (al,ar), left1) <- trial id as bs,
      (partition -> (bl,br), left2) <- trial flip bs as = case (left1,left2) of
            ([], []) -> C ar -- or br they should be similar?
            (_ , []) -> R bl
            ([], _ ) -> L al

data Zip f a b = C (f (a,b)) | L (f a) | R (f b)


deriving instance (Show (f (a,b)), Show (f a), Show (f b)) => Show (Zip f a b)

trial flip xs ys =
    T.mapM (\x -> do
            yys <- get
            case yys of
                [] -> do
                    return (Left x)
                y:ys -> do
                    put ys
                    return $ Right (flip (,) x y))
        xs `runState` F.toList ys

data NonEmptyList a = NonEmptyList a [a] deriving (Eq, Show)
data ZipAlign a b = Align [(a, b)] | RestA (NonEmptyList a) | RestB (NonEmptyList b) deriving (Eq, Show)

zipAlign :: [a] -> [b] -> ZipAlign a b
zipAlign [] [] = Align []
zipAlign (h:t) [] = RestA (NonEmptyList h t)
zipAlign [] (h:t) = RestB (NonEmptyList h t)
zipAlign (h:t) (h':t') =
    case zipAlign t t' of
        RestA r -> RestA r
        RestB r -> RestB r
        Align a -> Align ((h,h') : a)

{-

*Main> zipAlignG [1,2,3,4] [1,2,3,4,5]
R [5]

*Main> zipAlign [1,2,3,4] [1,2,3,4,5]
RestB (NonEmptyList 5 [])

*Main> zipAlignG2 [1,2,3,4] [1,2,3,4,5]
R [Nothing,Nothing,Nothing,Nothing,Just 5]



*Main> zipAlignG2 [1,2] [2,1]
C [(Just 1,Just 2),(Just 2,Just 1)]

*Main> zipAlignG [1,2] [2,1]
C [(1,2),(2,1)]

*Main> zipAlign [1,2] [2,1]
Align [(1,2),(2,1)]


-}

line 36: missing Ord k =>

M.partition :: (a -> Bool) -> M.Map k a -> (M.Map k a, M.Map k a) -- doesn't need any comparisons because the tree is already in order