FingerTrees.hs

{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE RecordWildCards #-}
module FingerTrees where

import Control.Lens hiding (Empty, (<|), (|>), deep)
import Data.Foldable (fold)
import Data.Monoid
import Data.Text (Text)
import qualified Data.Text as T
import GHC.Exts hiding (One)
import GHC.Generics (Generic)

data Affix a = One a
             | Two a a
             | Three a a a
             | Four a a a a
             deriving (Show, Eq, Ord, Functor, Foldable, Traversable)

instance IsList (Affix a) where
  type Item (Affix a) = a

  toList = foldMap pure

  fromList [a] = One a
  fromList [a, b] = Two a b
  fromList [a, b, c] = Three a b c
  fromList [a, b, c, d] = Four a b c d
  fromList _ = error "Affix must contain 1 to 4 elements"

-- 'v' is the type of the annotation.
data Node v a = Branch3 v a a a
              | Branch2 v a a
              deriving (Show, Eq, Ord, Functor, Foldable, Traversable)

instance Measured a v => IsList (Node v a) where
  type Item (Node v a) = a

  toList (Branch2 _ x y) = [x, y]
  toList (Branch3 _ x y z) = [x, y, z]

  fromList [x, y] =
    Branch2 (measure x <> measure y) x y
  fromList [x, y, z] =
    Branch3 (measure x <> measure y <> measure z) x y z
  fromList _ = error "Node must contain two or three elements"

data FingerTree v a
  = Empty
  | Single a
  | Deep {
    annotation :: v, -- Add an annotation to each branch.
    prefix :: Affix a,
    deeper :: FingerTree v (Node v a),
    suffix :: Affix a
  }
  deriving (Show, Eq, Ord, Functor, Foldable, Traversable)

instance Measured a v => Semigroup (FingerTree v a) where
  left <> right = case viewr left of
    Nil -> right
    View x left' -> left' <> (x <| right)

instance Measured a v => Monoid (FingerTree v a) where
  mempty = Empty

class Monoid v => Measured a v where
  measure :: a -> v

instance Measured a v => Measured (FingerTree v a) v where
  measure Empty = mempty
  measure (Single x) = measure x
  measure tree = annotation tree

instance Measured a v => Measured (Node v a) v where
  measure (Branch2 v _ _) = v
  measure (Branch3 v _ _ _) = v

instance Measured a v => Measured (Affix a) v where
  measure (One x) = measure x
  measure (Two x y) = foldMap measure [x, y]
  measure (Three x y z) = foldMap measure [x, y, z]
  measure (Four x y z q) = foldMap measure [x, y, z, q]

---------------
--- Helpers ---
---------------

-- Convert an affix into an entire tree, doing rebalancing if necessary.
affixToTree :: Measured a v => Affix a -> FingerTree v a
affixToTree affix = case affix of
  One a        -> Single a
  Two a b      -> Deep (measure affix) (One a) Empty (One b)
  Three a b c  -> Deep (measure affix) (One a) Empty (Two b c)
  Four a b c d -> Deep (measure affix) (Two a b) Empty (Two c d)


-- The `deep` function creates `Deep` finger trees.
deep' :: Measured a v => [a] -> FingerTree v (Node v a) -> [a] -> FingerTree v a
deep' prefix deeper suffix = case (prefix, suffix) of
  ([], []) -> case viewl deeper of
    Nil -> Empty
    View node deeper' -> deep' (toList node) deeper' []
  ([], _) -> case viewr deeper of
    Nil -> affixToTree $ fromList suffix
    View node deeper' -> deep' (toList node) deeper' suffix
  (_, []) -> case viewr deeper of
    Nil -> affixToTree $ fromList prefix
    View node deeper' -> deep' prefix deeper' (toList node)
  _ -> if length prefix > 4 || length suffix > 4
          then error "Affixes can be no more then length 4"
          else Deep annotation (fromList prefix) deeper (fromList suffix)
  where
    annotation = foldMap measure prefix <> measure deeper <> foldMap measure suffix

deep :: Measured a v => Affix a -> FingerTree v (Node v a) -> Affix a -> FingerTree v a
deep prefix deeper suffix =
  Deep (measure prefix <> measure deeper <> measure suffix) prefix deeper suffix

------------------------
--- Append / Prepend ---
------------------------

affixPrepend :: a -> Affix a -> Affix a
affixPrepend a = fromList . (a :) . toList

affixAppend :: a -> Affix a -> Affix a
affixAppend x = fromList . (++ [x]) . toList

infixr 5 <|

(<|) :: Measured a v => a -> FingerTree v a -> FingerTree v a
x <| Empty = Single x
x <| Single y = deep (One x) Empty (One y)
x <| Deep annot prefix deeper suffix = case prefix of
  Four a b c d -> Deep annot' (Two x a) (fromList [b, c, d] <| deeper) suffix
  prefix -> Deep annot' (affixPrepend x prefix) deeper suffix
  where
    annot' = measure x <> annot

infixl 5 |>

(|>) :: Measured a v => FingerTree v a -> a -> FingerTree v a
Empty |> y = Single y
Single x |> y = Deep (measure x <> measure y) (One x) Empty (One y)
Deep annot prefix deeper (Four a b c d) |> y = Deep annot' prefix (deeper |> node) (Two d y)
  where
    annot' = measure y <> annot
    node = Branch3 (foldMap measure [a, b, c]) a b c
tree |> y = tree { suffix = affixAppend y $ suffix tree }

-------------
--- Views ---
-------------

data View v a = Nil | View a (FingerTree v a) deriving Show

-- Shows the first element and the remaining tree
viewl :: Measured a v => FingerTree v a -> View v a
viewl Empty = Nil
viewl (Single x) = View x Empty
viewl (Deep _ (One x) deeper suffix) =
  case viewl deeper of
    View node rest' ->
      let pref = fromList $ toList node
          annot = measure pref <> measure rest' <> measure suffix
      in View x $ Deep annot pref rest' suffix
    Nil -> View x $ affixToTree suffix
viewl (Deep _ prefix deeper suffix) =
  let first:rest = toList prefix
      prefix' = fromList rest
      annot = measure prefix' <> measure deeper <> measure suffix
  in View first $ Deep annot prefix' deeper suffix

-- Shows the last element and the preceding tree
viewr :: Measured a v => FingerTree v a -> View v a
viewr Empty = Nil
viewr (Single x) = View x Empty
viewr (Deep _ prefix deeper (One x)) =
  case viewr deeper of
    View node rest' ->
      let suff = fromList $ toList node
          annot = measure prefix <> measure rest' <> measure suff
      in View x $ Deep annot prefix rest' suff
    Nil -> View x $ affixToTree prefix
viewr (Deep _ prefix deeper suffix) =
  let annot = measure prefix <> measure deeper <> measure  suffixInit
      suffixLast = last $ toList suffix
      suffixInit = fromList $ init $ toList suffix
  in View suffixLast $ Deep annot prefix deeper suffixInit

---------------
--- Lookups ---
---------------

splitList
  :: Measured a v
  => (v -> Bool) -- Monotonic predicate on annotations.
  -> v           -- Left-most annotation.
  -> [a]         -- List of measurable values.
  -> ([a], [a])
splitList pred start [] = error "Split point not found"
splitList pred start (x:xs) =
  let start' = start <> measure x -- The new left-most annotation value after including this element.
  in if pred start'
       then ([], x:xs)
       else let (before, after) = splitList pred start' xs
            in (x : before, after)

listExample :: ([Value Char], [Value Char])
listExample = splitList (\(Size s) -> s > 5) (Size 0) (Value <$> ['a' .. 'z'])

-- A Zipper for a FingerTree
data Split t a = Split t a t deriving Show
type SplitTree v a = Split (FingerTree v a) a

split :: Measured a v
      => (v -> Bool)     -- Monotonic predicate on annotations.
      -> v               -- Annotation on the left end of the subsequence.
      -> FingerTree v a  -- Subsequence to search within.
      -> SplitTree v a
split _ _ Empty = error "Split point not found"
split pred start (Single x) =
  if pred (start <> measure x)
    then Split Empty x Empty
    else error "Split point not found"
split pred v (Deep total prefix deeper suffix)
  | not (pred $ v <> total) = error "Split point not found"
  -- Split is in the prefix
  | pred (v <> measure prefix) =
    let (before, x:after) = splitList pred v (toList prefix)
    in Split (chunkToTree before) x (deep (fromList after) deeper suffix)
  -- Split is in the deeper tree
  | pred (v <> measure prefix <> measure deeper) =
    let Split before node after = split pred (v <> measure prefix) deeper
        (beforeNode, x:afterNode) = splitList pred (v <> measure prefix <> measure before) (toList node)
    in Split (deep prefix before (fromList beforeNode)) x (deep (fromList afterNode) after suffix)
  -- Split is in the suffix
  | otherwise =
    let (before, x:after) = splitList pred (v <> measure prefix <> measure deeper) (toList suffix)
    in Split (deep prefix deeper (fromList before)) x (chunkToTree after)
  where
    chunkToTree [] = Empty
    chunkToTree xs = affixToTree $ fromList xs

data SearchResult v a =
    Position (FingerTree v a) a (FingerTree v a)
  | OnLeft
  | OnRight
  | Nowhere
  deriving (Show, Eq, Ord, Generic)

search :: Measured a v => (v -> v -> Bool) -> FingerTree v a -> SearchResult v a
search p t
  | pLeft && pRight = OnLeft
  | not pLeft && pRight = let Split l x r = searchTree p mempty t mempty in Position l x r
  | not pLeft && not pRight = OnRight
  | otherwise = Nowhere
  where
    pLeft = p mempty (measure t)
    pRight = p (measure t) mempty

searchTree :: Measured a v =>
    (v -> v -> Bool) -> v -> FingerTree v a -> v -> Split (FingerTree v a) a
searchTree p vl Empty vr = error "Cannot search an empty tree"
searchTree p _ (Single a) _ = Split Empty a Empty
searchTree p vl (Deep v prefix deeper suffix) vr
  | p vlPrefix deepVrSuf = let Split l x r = searchDigit p vl prefix deepVrSuf
                           in Split (maybe Empty affixToTree l) x (deepL r deeper suffix)
  | p vlPreDeep vrSuffix = let Split ml x mr = searchTree p vlPrefix deeper vrSuffix
                               Split l x' r = searchNode p (vlPrefix <> measure ml) x (measure mr <> vrSuffix)
                           in Split (deepR prefix ml l) x' (deepL r mr suffix)
  | otherwise            = let Split l x r = searchDigit p vlPreDeep suffix vr
                           in Split (deepR prefix deeper l) x (maybe Empty affixToTree r)
  where
    vlPrefix = vl <> measure prefix
    vlPreDeep = vlPrefix <> measure deeper
    vrSuffix = measure suffix <> vr
    deepVrSuf = measure deeper <> vrSuffix

rotL :: Measured a v => FingerTree v (Node v a) -> Affix a -> FingerTree v a
rotL tree affix = case viewl tree of
  Nil -> affixToTree affix
  View a tree' -> Deep (measure tree <> measure affix) (fromList $ toList a) tree' affix

rotR :: Measured a v => Affix a -> FingerTree v (Node v a) -> FingerTree v a
rotR prefix tree = case viewr tree of
  Nil -> affixToTree prefix
  View a tree' -> Deep (measure prefix <> measure tree) prefix tree' (fromList $ toList a)

deepL :: Measured a v => Maybe (Affix a) -> FingerTree v (Node v a) -> Affix a -> FingerTree v a
deepL Nothing m sf   = rotL m sf
deepL (Just pr) m sf = deep pr m sf

deepR :: Measured a v => Affix a -> FingerTree v (Node v a) -> Maybe (Affix a) -> FingerTree v a
deepR pr m Nothing   = rotR pr m
deepR pr m (Just sf) = deep pr m sf

searchNode :: Measured a v => (v -> v -> Bool) -> v -> Node v a -> v -> Split (Maybe (Affix a)) a
searchNode p vl (Branch3 _ a b c) vr
  | p vla vrbc = Split Nothing a (Just (Two b c))
  | p vlab vrc = Split (Just (One a)) b (Just (One c))
  | otherwise = Split (Just (Two a b)) c Nothing
  where
    vla = vl <> measure a
    vlab = vla <> measure b
    vrc = measure c <> vr
    vrbc = measure b <> vrc
searchNode p vl (Branch2 _ a b) vr
  | p (vl <> measure a) (measure b <> vr) = Split Nothing a (Just (One b))
  | otherwise = Split (Just (One a)) b Nothing

searchDigit :: Measured a v => (v -> v -> Bool) -> v -> Affix a -> v -> Split (Maybe (Affix a)) a
searchDigit p vl (One a) vr = vl `seq` vr `seq` Split Nothing a Nothing
searchDigit p vl (Two a b) vr
  | p (vl <> measure a) (measure b <> vr) = Split Nothing a (Just $ One b)
  | otherwise = Split (Just $ One a) b Nothing
searchDigit p vl (Three a b c) vr
  | p (vl <> measure a) (measure b <> measure c <> vr) = Split Nothing a (Just $ Two b c)
  | p (vl <> measure a <> measure b) (measure c <> vr) = Split (Just $ One a) b (Just $ One c)
  | otherwise = Split (Just $ Two a b) c Nothing
searchDigit p vl (Four a b c d) vr
  | p (vl <> measure a) (measure b <> measure c <> measure d <> vr) = Split Nothing a (Just $ Three b c d)
  | p (vl <> measure a <> measure b) (measure c <> measure d <> vr) = Split (Just $ One a) b (Just $ Two c d)
  | p (vl <> measure a <> measure b <> measure c) (measure d <> vr) = Split (Just $ Two a b) c (Just $ One d)
  | otherwise = Split (Just $ Three a b c) d Nothing

----------------------
--- Small Examples ---
----------------------

oneChunk :: FingerTree (Sum Int) String
oneChunk = Single "Foo"

threeChunk :: FingerTree (Sum Int) String
threeChunk = Deep (Sum 3) (One "Foo") Empty (Two "Bar" "Baz")


deeperTree :: FingerTree (Sum Int) String
deeperTree =
  Deep
    { annotation = Sum 3
    , prefix = One "Foo"
    , deeper = Single (Branch2 (Sum 3) "Baz" "Qux")
    , suffix = One "Bar"
    }

evenDeeperTree :: FingerTree (Sum Int) String
evenDeeperTree =
  Deep
    { annotation = Sum 3
    , prefix = Two "This" "is"
    , deeper = Deep
                 { annotation = Sum 3
                 , prefix = One (Branch2 (Sum 3) "Bar" "Baz")
                 , deeper = Empty
                 , suffix = One (Branch2 (Sum 3) "Qux" "Bop")
                 }
    , suffix = One "Beep"
    }

------------------------
--- Complete Example ---
------------------------

newtype Size = Size Int
 deriving (Show, Eq, Ord)

instance Semigroup Size where
  (Size x) <> (Size y) = Size (x + y)

instance Monoid Size where
  mempty = Size 0

newtype Value a = Value a
  deriving Show

instance Measured (Value a) Size where
  measure (Value _) = Size 1

mkBranch2 :: Measured a v => a -> a -> Node v a
mkBranch2 x y = Branch2 (measure x <> measure y) x y

mkBranch3 :: Measured a v => a -> a -> a -> Node v a
mkBranch3 x y z = Branch3 (foldMap measure [x, y, z]) x y z

layer3 :: FingerTree Size (Node Size (Node Size (Value Char)))
layer3 = Empty

layer2 :: FingerTree Size (Node Size (Value Char))
layer2 = Deep annotation prefix layer3 suffix
  where
    annotation = measure prefix <> measure layer3 <> measure suffix
    prefix = Two (mkBranch2 (Value 'i') (Value 's')) (mkBranch2 (Value 'i') (Value 's'))
    suffix = Two (mkBranch3 (Value 'n') (Value 'o') (Value 't')) (mkBranch2 (Value 'a') (Value 't'))

layer1 :: FingerTree Size (Value Char)
layer1 = Deep annotation prefix layer2 suffix
  where
    annotation = measure prefix <> measure layer2 <> measure suffix
    prefix = Two (Value 't') (Value 'h')
    suffix = Three (Value 'r') (Value 'e') (Value 'e')

exampleTree :: FingerTree Size (Value Char)
exampleTree = layer1

exampleSplit :: SplitTree Size (Value Char)
exampleSplit = split (\(Size i) -> i > 5) mempty exampleTree

----------------------
--- Priority Queue ---
----------------------

data Prioritized a = Prioritized { priority :: Int, item :: a } deriving Show
data Priority = NegInf | Priority Int deriving (Show, Eq, Ord)

instance Semigroup Priority where
  (<>) = max

instance Monoid Priority where
  mempty = NegInf

instance Measured (Prioritized a) Priority where
  measure = Priority . priority

newtype PriorityQueue a = PriorityQueue (FingerTree Priority (Prioritized a))
  deriving Show

push :: PriorityQueue a -> a -> Int -> PriorityQueue a
push (PriorityQueue tree) a p = PriorityQueue $ Prioritized p a <| tree

pop :: PriorityQueue a -> (a, PriorityQueue a)
pop (PriorityQueue tree) =
  let maxPriority = measure tree
      Split left x right = split (== maxPriority) mempty tree
  in (item x, PriorityQueue $ left <> right)

emptyQueue :: PriorityQueue String
emptyQueue = PriorityQueue Empty

strings :: [String]
strings = ["Goodbye", "one", "Hello", "a"]

queue = foldr addToQueue emptyQueue strings
  where addToQueue str queue = push queue str (length str)

------------
--- Rope ---
------------

newtype Width = Width { unWidth :: Int}
  deriving (Eq, Ord, Show, Num) via Int
  deriving (Generic)

instance Semigroup Width where
  (<>) = (+)

instance Monoid Width where
  mempty = 0

instance Measured Text Width where
  measure :: Text -> Width
  measure piece = Width (T.length piece)

newtype Rope = Rope { unRope :: FingerTree Width Text }
  deriving (Generic)

instance Semigroup Rope where
  Rope a <> Rope b = Rope $ a <> b

instance Monoid Rope where
  mempty = Rope Empty

instance Show Rope where
  show text = "\"" ++ T.unpack (fromRope text) ++ "\""

emptyRope :: Rope
emptyRope = mempty

singletonRope :: Char -> Rope
singletonRope = Rope . Single . T.singleton

replicateRope :: Int -> Rope -> Rope
replicateRope i (Rope tree) =
  Rope $ foldr (\_ acc -> tree <> acc) mempty [1 .. i]

replicateChar :: Int -> Char -> Rope
replicateChar i = Rope . Single . T.replicate i . T.singleton

widthRope :: Rope -> Int
widthRope = unWidth . measure . unRope

splitRope :: Int -> Rope -> (Rope, Rope)
splitRope i (Rope tree) =
  case search (\w1 _ -> w1 >= Width i) tree of
    Position before a after ->
      let Width w = measure before
          (one, two) = T.splitAt (i - w) a
      in (Rope $ before |> one, Rope $ two <| after)
    OnLeft -> (Rope Empty, Rope tree)
    OnRight -> (Rope tree, Rope Empty)
    Nowhere -> error "Out of bounds index"

insertRope :: Int -> Rope -> Rope -> Rope
insertRope i (Rope new) old =
  let (Rope before, Rope after) = splitRope i old
  in Rope $ before <> new <> after

findIndexRope :: (Char -> Bool) -> Rope -> Maybe Int
findIndexRope p (Rope tree) = fst $ foldl f (Nothing, 0) tree
  where
    f :: (Maybe Int, Int) -> Text -> (Maybe Int, Int)
    f acc next = case acc of
      (Just i, _) -> (Just i, 0)
      (Nothing, i) -> case T.findIndex p next of
        Nothing -> (Nothing, i + T.length next)
        Just j -> (Just (i + j), 0)

appendRope :: Text -> Rope -> Rope
appendRope t (Rope tree) = Rope $ t <| tree

fromRope :: Rope -> Text
fromRope = foldr T.append mempty . unRope

intoRope :: Text -> Rope
intoRope = Rope . Single

nullRope :: Rope -> Bool
nullRope (Rope tree) = case viewl tree of
  Nil -> False
  View x _ -> T.null x

splitRange :: Int -> Int -> Rope -> Split Rope Rope
splitRange i j rope =
  let (before, inter) = splitRope i rope
      (target, after) = splitRope (j - i) inter
  in Split before target after

----------------------
--- Annotated Rope ---
----------------------

newtype Rope' = Rope' { unRope' :: FingerTree Width Chunk }
  deriving (Generic)

-- .. | Italic | H1 | H2 | H3 | H4 | H5 | H6
data Annotation = Bold
  deriving (Eq, Ord, Show)

data Chunk = Chunk { _chunkLength :: Int, _fromChunk :: Text, _chunkAnno :: Maybe Annotation }
  deriving (Eq, Ord, Show)
makeLenses ''Chunk

mkChunk :: T.Text -> Maybe Annotation -> Chunk
mkChunk txt = Chunk (T.length txt) txt

splitAtChunk :: Int -> Chunk -> (Chunk, Chunk)
splitAtChunk i Chunk{..} =
  let (before, after) = T.splitAt i _fromChunk
  in (mkChunk before _chunkAnno, mkChunk after _chunkAnno)

instance Measured Chunk Width where
  measure :: Chunk -> Width
  measure chunk = Width (_chunkLength chunk)

instance Semigroup Rope' where
  Rope' a <> Rope' b = Rope' $ a <> b

instance Monoid Rope' where
  mempty = Rope' Empty

instance Show Rope' where
  show text = "\"" ++ T.unpack (fromRope' text) ++ "\""

singletonRope' :: Char -> Rope'
singletonRope' = Rope' . Single . flip mkChunk Nothing . T.singleton

replicateRope' :: Int -> Rope' -> Rope'
replicateRope' i (Rope' tree) =
  Rope' $ foldr (\_ acc -> tree <> acc) mempty [1 .. i]

replicateChar' :: Int -> Char -> Rope'
replicateChar' i = Rope' . Single . flip mkChunk Nothing . T.replicate i . T.singleton

widthRope' :: Rope' -> Int
widthRope' = unWidth . measure . unRope'

splitRope' :: Int -> Rope' -> (Rope', Rope')
splitRope' i (Rope' tree) =
  case search (\w1 _ -> w1 >= Width i) tree of
    Position before a after ->
      let Width w = measure before
          (one, two) = splitAtChunk (i - w) a
      in (Rope' $ before |> one, Rope' $ two <| after)
    OnLeft -> (Rope' Empty, Rope' tree)
    OnRight -> (Rope' tree, Rope' Empty)
    Nowhere -> error "Out of bounds index"

insertRope' :: Int -> Rope' -> Rope' -> Rope'
insertRope' i (Rope' new) old =
  let (Rope' before, Rope' after) = splitRope' i old
  in Rope' $ before <> new <> after

findIndexRope' :: (Char -> Bool) -> Rope' -> Maybe Int
findIndexRope' p (Rope' tree) = fst $ foldl f (Nothing, 0) tree
  where
    f :: (Maybe Int, Int) -> Chunk -> (Maybe Int, Int)
    f acc next = case acc of
      (Just i, _) -> (Just i, 0)
      (Nothing, i) -> case T.findIndex p $ _fromChunk next of
        Nothing -> (Nothing, i + T.length (_fromChunk next))
        Just j -> (Just (i + j), 0)

appendRope' :: Chunk -> Rope' -> Rope'
appendRope' t (Rope' tree) = Rope' $ t <| tree

fromRope' :: Rope' -> Text
fromRope' = foldr (T.append . T.pack . show) mempty . unRope'

intoRope' :: Text -> Rope'
intoRope' = Rope' . Single . flip mkChunk Nothing

nullRope' :: Rope' -> Bool
nullRope' (Rope' tree) = case viewl tree of
  Nil -> False
  View x _ -> T.null $ _fromChunk x

splitRange' :: Int -> Int -> Rope' -> Split Rope' Rope'
splitRange' i j rope =
  let (before, inter) = splitRope' i rope
      (target, after) = splitRope' (j - i) inter
  in Split before target after

applyAnnoToRange :: Maybe Annotation -> Int -> Int -> Rope' -> Rope'
applyAnnoToRange anno i j rope =
  let Split before (Rope' x) after = splitRange' i j rope
      x' = fmap (\chunk -> chunk & chunkAnno .~ anno) x
  in before <> Rope' x' <> after