Attempt at type-aligned finger trees

Aligned.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}

{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module FingerTree.Aligned where

import Control.Category
import Prelude hiding ((.), id)

import qualified Aligned.Internal as A

newtype NT g a b c = NT {runNT :: g a b -> g a c}

instance Category (NT g a) where
  id = NT id
  NT f . NT g = NT (f . g)

class AFoldable t where
  afoldMap :: Category c => (forall a b. f a b -> c a b) -> t f a b -> c a b
  afoldr :: (forall a b c. f b c -> g a b -> g a c) -> g a b -> t f b c -> g a c
  afoldMap f = afoldr ((.) . f) id
  afoldr f b t = runNT (afoldMap (NT . f) t) b
  afoldl :: (forall a b c. g b c -> f a b -> g a c) -> g b c -> t f a b -> g a c
  afoldl f b t = A.runOp $ runNT (A.runOp $ afoldMap (\fab -> A.Op (NT (A.Op . flip f fab . A.runOp))) t) (A.Op b)
  toThrist :: t f a b -> A.Thrist f a b
  toThrist = afoldMap (A.:. A.Id)
  {-# MINIMAL afoldMap | afoldr #-}

consMany :: (AFoldable t, A.Cons t') => t f b c -> t' f a b -> t' f a c
consMany xs ys = afoldr A.cons ys xs

snocMany :: (A.Snoc t, AFoldable t') => t f b c -> t' f a b -> t f a c
snocMany xs ys = afoldl A.snoc xs ys

instance AFoldable t => AFoldable (A.Rev t) where
  afoldMap f = A.runOp . afoldMap (A.Op . f . A.runOp) . A.runRev

instance AFoldable A.Thrist where
  afoldMap _ A.Id = id
  afoldMap f (x A.:. xs) = f x . afoldMap f xs

instance AFoldable A.Cat where
  afoldMap = A.foldCat

class AFunctor t where
  afmap :: (forall a b. f a b -> g a b) -> t f a b -> t g a b

instance AFunctor t => AFunctor (A.Rev t) where
  afmap f = A.Rev . afmap (A.Op . f . A.runOp) . A.runRev

instance AFunctor A.Thrist where
  afmap _ A.Id = A.Id
  afmap f (x A.:. xs) = f x A.:. afmap f xs

instance AFunctor A.Q where
  afmap g (A.Q f r s) = A.Q (afmap g f) (afmap g r) (afmap g s)

instance AFunctor A.Cat where
  afmap _ A.E = A.E
  afmap f (A.C x xs) = A.C (f x) (afmap (afmap f) xs)

data Node f a b where
  Node2 :: f b c -> f a b -> Node f a c
  Node3 :: f c d -> f b c -> f a b -> Node f a d

instance AFoldable Node where
  afoldMap f (Node2 p q) = f p . f q
  afoldMap f (Node3 p q r) = f p . f q . f r

instance AFunctor Node where
  afmap f (Node2 p q) = Node2 (f p) (f q)
  afmap f (Node3 p q r) = Node3 (f p) (f q) (f r)

data Digit f a b where
  One :: f a b-> Digit f a b
  Two :: f b c -> f a b -> Digit f a c
  Three :: f c d -> f b c -> f a b -> Digit f a d
  Four :: f d e -> f c d -> f b c -> f a b -> Digit f a e

instance AFoldable Digit where
  afoldMap f (One p) = f p
  afoldMap f (Two p q) = f p . f q
  afoldMap f (Three p q r) = f p . f q . f r
  afoldMap f (Four p q  r s) = f p . f q . f r . f s

instance AFunctor Digit where
  afmap f (One p) = One (f p)
  afmap f (Two p q) = Two (f p) (f q)
  afmap f (Three p q r) = Three (f p) (f q) (f r)
  afmap f (Four p q r s) = Four (f p) (f q) (f r) (f s)

nodeToDigit :: Node f a b -> Digit f a b
nodeToDigit (Node2 p q) = Two p q
nodeToDigit (Node3 p q r) = Three p q r

data FingerTree f a b where
  Empty :: FingerTree f a a
  Single :: f a b -> FingerTree f a b
  Deep :: Digit f c d -> FingerTree (Node f) b c -> Digit f a b -> FingerTree f a d

instance AFoldable FingerTree where
  afoldMap _ Empty = id
  afoldMap f (Single p) = f p
  afoldMap f (Deep a b c) = afoldMap f a . afoldMap (afoldMap f) b . afoldMap f c

instance AFunctor FingerTree where
  afmap _ Empty = Empty
  afmap f (Single p) = Single (f p)
  afmap f (Deep pr m sf) = Deep (afmap f pr) (afmap (afmap f) m) (afmap f sf)

digitToTree :: Digit f a b -> FingerTree f a b
digitToTree (One p) = Single p
digitToTree (Two p q) = Deep (One p) Empty (One q)
digitToTree (Three p q r) = Deep (Two p q) Empty (One r)
digitToTree (Four p q r s) = Deep (Two p q) Empty (Two r s)

instance A.Nil FingerTree where
  nil = Empty

instance A.Singleton FingerTree where
  singleton = Single

instance A.Cons FingerTree where
  p `cons` Empty = Single p
  p `cons` Single q = Deep (One p) Empty (One q)
  p `cons` Deep (Four q r s t) m sf = Deep (Two p q) (Node3 r s t `A.cons` m) sf
  p `cons` Deep (Three q r s) m sf = Deep (Four p q r s) m sf
  p `cons` Deep (Two q r) m sf = Deep (Three p q r) m sf
  p `cons` Deep (One q) m sf = Deep (Two p q) m sf

instance A.Snoc FingerTree where
  Empty `snoc` p = Single p
  Single p `snoc` q = Deep (One p) Empty (One q)
  Deep pr m (Four p q r s) `snoc` t = Deep pr (m `A.snoc` Node3 p q r) (Two s t)
  Deep pr m (Three p q r) `snoc` s = Deep pr m (Four p q r s)
  Deep pr m (Two p q) `snoc` r = Deep pr m (Three p q r)
  Deep pr m (One p) `snoc` q = Deep pr m (Two p q)

instance A.Uncons FingerTree where
  uncons Empty = A.Empty
  uncons (Single p) = p A.:&: Empty
  uncons (Deep (One p) m sf) = case A.uncons m of
    A.Empty -> p A.:&: digitToTree sf
    q A.:&: m' -> p A.:&: Deep (nodeToDigit q) m' sf
  uncons (Deep (Two p q) m sf) = p A.:&: Deep (One q) m sf
  uncons (Deep (Three p q r) m sf) = p A.:&: Deep (Two q r) m sf
  uncons (Deep (Four p q r s) m sf) = p A.:&: Deep (Three q r s) m sf

instance A.Unsnoc FingerTree where
  unsnoc Empty = A.Empty
  unsnoc (Single p) = Empty A.:&: p
  unsnoc (Deep pr m (One p)) = case A.unsnoc m of
    A.Empty -> digitToTree pr A.:&: p
    m' A.:&: q -> Deep pr m' (nodeToDigit q) A.:&: p
  unsnoc (Deep pr m (Two p q)) = Deep pr m (One p) A.:&: q
  unsnoc (Deep pr m (Three p q r)) = Deep pr m (Two p q) A.:&: r
  unsnoc (Deep pr m (Four p q r s)) = Deep pr m (Three p q r) A.:&: s

app3 :: FingerTree f c d -> A.Thrist f b c -> FingerTree f a b -> FingerTree f a d
app3 Empty ts xs = consMany ts xs
app3 xs ts Empty = snocMany xs ts
app3 (Single x) ts xs = A.cons x (consMany ts xs)
app3 xs ts (Single x) = A.snoc (snocMany xs ts) x
app3 (Deep pr1 m1 sf1) ts  (Deep pr2 m2 sf2) = Deep pr1 (app3 m1 (nodes (toThrist sf1 . ts . toThrist pr2)) m2) sf2
  where
    nodes :: A.Thrist f a b -> A.Thrist (Node f) a b
    nodes A.Id = error "Unreachable: app3.nodes"
    nodes (_ A.:. A.Id) = error "Unreachable: app3.nodes"
    nodes (a A.:. b A.:. A.Id) = Node2 a b A.:. A.Id
    nodes (a A.:. b A.:. c A.:. A.Id) = Node3 a b c A.:. A.Id
    nodes (a A.:. b A.:. c A.:. d A.:. A.Id) = Node2 a b A.:. Node2 c d A.:. A.Id
    nodes (a A.:. b A.:. c A.:. xs) = Node3 a b c A.:. nodes xs

instance Category (FingerTree f) where
  id = Empty
  xs . ys = app3 xs A.Id ys