Implementing weight-balanced trees using recursion schemes. -- https://stackoverflow.com/questions/74651343/using-a-paramorphism-inside-of-an-apomorphism

A.hs

-- Implementing weight-balanced trees using recursion schemes.
-- (actually just concat3 from Adams's report)
-- https://stackoverflow.com/questions/74651343/using-a-paramorphism-inside-of-an-apomorphism
--
-- References for WBT:
-- - Implementing Sets Efficiently in a Functional Language, Stephen Adams, 1992
--   https://ia600700.us.archive.org/0/items/djoyner-papers/SHA256E-s234215--592b97774eca4193a05ed9472ab6e23788d3a0bea5d1b98cef301460ab4010ee.pdf
-- - Binary search trees of bounded balance, J. Nievergelt and E. M. Reingold
--   https://dl.acm.org/doi/10.1145/800152.804906
-- - Balancing Weight-Balanced Trees, Y. Hirai and K. Yamamoto
--   https://yoichihirai.com/bst.pdf

{-# LANGUAGE DeriveFunctor, UndecidableInstances, FlexibleContexts #-}

-- Fixed point of a Functor
newtype Fix f = In (f (Fix f))

deriving instance (Eq (f (Fix f))) => Eq (Fix f)
deriving instance (Ord (f (Fix f))) => Ord (Fix f)
deriving instance (Show (f (Fix f))) => Show (Fix f)

out :: Fix f -> f (Fix f)
out (In f) = f

type RAlgebra f a = f (Fix f, a) -> a

para :: (Functor f) => RAlgebra f a -> Fix f -> a
para rAlg = rAlg . fmap fanout . out
        where fanout t = (t, para rAlg t)

-- Apomorphism
type RCoalgebra f a = a -> f (Either (Fix f) a)

apo :: Functor f => RCoalgebra f a -> a -> Fix f
apo rCoalg = In . fmap fanin . rCoalg
        where fanin = either id (apo rCoalg)

data TreeF a x = E | T a Int x x
  deriving Functor

type Tree a = Fix (TreeF a)

-- | Smart constructor that computes the weight.
_N :: a -> Tree a -> Tree a -> Tree a
_N v l r = In (T v (1 + weight l + weight r) l r)
  where
    weight (In E) = 0
    weight (In (T _ w _ _)) = w

balance :: Int
balance = 4

add :: Tree a -> a -> Tree a
add = error "todo"

-- | Tree with delayed rebalancing operation T'.
data TreeF1 a x = E1 | T' a x x | Id (Tree a)
  deriving Functor

concatAlg :: Ord a => a -> RCoalgebra (TreeF1 a) (Tree a, Tree a)
concatAlg v (In E, r) = Id (add r v)
concatAlg v (l, In E) = Id (add l v)
concatAlg v (l@(In (T v1 n1 l1 r1)), r@(In (T v2 n2 l2 r2))) =
  if balance * n1 < n2 then T' v2 (Right (l, l2)) (Left (In (Id r2)))
  else if balance * n2 < n1 then T' v1 (Left (In (Id l1))) (Right (r1, r))
  else Id (_N v1 l r)

_T :: a -> Tree a -> Tree a -> Tree a
_T = error "todo: rebalance"

type Algebra f a = f a -> a

-- do the rebalancing on T' v l r nodes
rebalanceAlg :: Algebra (TreeF1 a) (Tree a)
rebalanceAlg E1 = In E
rebalanceAlg (T' v l r) = _T v l r
rebalanceAlg (Id t) = t

cata :: Functor f => Algebra f a -> Fix f -> a
cata alg (In x) = alg (cata alg <$> x)

concat3 :: Ord a => a -> Tree a -> Tree a -> Tree a
concat3 v l r = (cata rebalanceAlg . apo (concatAlg v)) (l, r)

cataApo :: Functor f => Algebra f b -> RCoalgebra f a -> a -> b
cataApo alg coalg = go
  where
    go x = alg (either (cata alg) go <$> coalg x)

concat3' :: Ord a => a -> Tree a -> Tree a -> Tree a
concat3' v l r = cataApo rebalanceAlg (concatAlg v) (l, r)