{-# LANGUAGE DeriveTraversable #-}

-- | A 'LexTree' is a map designed to reconstruct the lexical structure (tree of scopes) of a source file, given an unordered list of scopes.
-- Values are inserted with a pair source locations as its key.
-- For a given key, we can then ask what the smallest enclosing scope is.
--
-- For example, in the snippet below the smallest scope containing @x@ is @b@.
-- @
--      x
-- |      a      |
--    |  b  |
--                   |   c   |
-- @
--
-- Scopes are not allowed to overlap.
--
-- The purpose of this data structure is to find out what surrounding definition a certain use site belongs to.
module Calligraphy.Util.LexTree
  ( LexTree (..),
    TreeError (..),
    Calligraphy.Util.LexTree.lookup,
    lookupOuter,
    insert,
    emptyLexTree,
    foldLexTree,
    toForest,
    insertWith,
    height,
    toList,
    bin,
    shift,
  )
where

import Control.Applicative
import Data.Tree (Forest, Tree (..))

data LexTree p a
  = Tip
  | Bin
      {-# UNPACK #-} !Int
      -- ^ Height
      !(LexTree p a)
      -- ^ Scopes at the same level, left of this one
      !p
      -- ^ Left-hand bound of this scope (inclusive)
      a
      !(LexTree p a)
      -- ^ Children
      !p
      -- ^ Right-hand bound of this scope (exclusive)
      !(LexTree p a)
      -- ^ Scopes at the same level, right of this entry
  deriving (Show, Functor, Foldable, Traversable)

instance (Eq p, Eq a) => Eq (LexTree p a) where
  ta == tb = toList ta == toList tb

lookup :: Ord p => p -> LexTree p a -> Maybe a
lookup p = foldLexTree Nothing f
  where
    f ls l a m r rs
      | p >= l && p < r = m <|> Just a
      | p < l = ls
      | p >= r = rs
      | otherwise = error "impossible"

lookupOuter :: Ord p => p -> LexTree p a -> Maybe a
lookupOuter p = foldLexTree Nothing f
  where
    f ls l a _ r rs
      | p >= l && p < r = Just a
      | p < l = ls
      | p >= r = rs
      | otherwise = error "impossible"

toList :: LexTree p a -> [(p, a, p)]
toList t = foldLexTree id f t []
  where
    f ls l a m r rs = ls . ((l, a, r) :) . m . rs

foldLexTree :: r -> (r -> p -> a -> r -> p -> r -> r) -> LexTree p a -> r
foldLexTree fTip fBin = go
  where
    go Tip = fTip
    go (Bin _ ls l a ms r rs) = fBin (go ls) l a (go ms) r (go rs)

emptyLexTree :: LexTree p a
emptyLexTree = Tip

toForest :: LexTree p a -> Forest (p, a, p)
toForest lt = foldLexTree id f lt []
  where
    f ls l a m r rs = ls . (Node (l, a, r) (m []) :) . rs

{-# INLINE height #-}
height :: LexTree p a -> Int
height Tip = 0
height (Bin h _ _ _ _ _ _) = h

shift :: Num p => p -> LexTree p a -> LexTree p a
shift p = go
  where
    go Tip = Tip
    go (Bin h ls l a ms r rs) = Bin h (go ls) (l + p) a (go ms) (r + p) (go rs)

data TreeError p a
  = -- | Nonsensical bounds, i.e. a left-hand bound larger than the right-hand obund
    InvalidBounds p a p
  | -- | Two identical scopes
    OverlappingBounds a a p p
  | -- | An attempt to split halfway through a scope, usually the result of two partially overlapping scopes
    MidSplit
  | -- | Attempting to insert a scope that would not form a tree structure
    LexicalError p a p (LexTree p a)
  deriving (Functor, Foldable, Traversable, Eq, Show)

{-# INLINE bin' #-}
bin' :: LexTree p a -> p -> a -> LexTree p a -> p -> LexTree p a -> LexTree p a
bin' ls l a ms r rs = Bin (max (height ls) (height rs) + 1) ls l a ms r rs

-- | Only works if the height difference of the two trees is at most 2
{-# INLINE bin #-}
bin :: LexTree p a -> p -> a -> LexTree p a -> p -> LexTree p a -> LexTree p a
bin (Bin lh lls ll la lms lr lrs) l a ms r rs
  | lh > height rs + 1 =
      case lrs of
        Bin lrh lrls lrl lra lrms lrr lrrs | lrh > height lls -> bin' (bin' lls ll la lms lr lrls) lrl lra lrms lrr (bin' lrrs l a ms r rs)
        _ -> bin' lls ll la lms lr (bin' lrs l a ms r rs)
bin ls l a ms r (Bin rh rls rl ra rms rr rrs)
  | rh > height ls + 1 =
      case rls of
        Bin rlh rlls rll rla rlms rlr rlrs | rlh > height rrs -> bin' (bin' ls l a ms r rlls) rll rla rlms rlr (bin' rlrs rl ra rms rr rrs)
        _ -> bin' (bin' ls l a ms r rls) rl ra rms rr rrs
bin ls l a ms r rs = bin' ls l a ms r rs

split :: Ord p => p -> LexTree p a -> Either (TreeError p a) (LexTree p a, LexTree p a)
split p = go
  where
    go Tip = pure (Tip, Tip)
    go (Bin _ ls l a ms r rs)
      | p <= l = do
          (ll, lr) <- go ls
          pure (ll, bin lr l a ms r rs)
      | p >= r = do
          (rl, rr) <- go rs
          pure (bin ls l a ms r rl, rr)
      | otherwise = Left MidSplit

{-# INLINE insertWith #-}
insertWith :: Ord p => (a -> a -> Maybe a) -> p -> a -> p -> LexTree p a -> Either (TreeError p a) (LexTree p a)
insertWith f il ia ir t
  | il >= ir = Left $ InvalidBounds il ia ir
  | otherwise = go t
  where
    go Tip = pure $ bin Tip il ia Tip ir Tip
    go (Bin h ls l a ms r rs)
      | ir <= l = flip fmap (go ls) $ \ls' -> bin ls' l a ms r rs
      | il >= r = flip fmap (go rs) $ \rs' -> bin ls l a ms r rs'
      | il == l && ir == r = case f ia a of
          Just a' -> pure $ Bin h ls l a' ms r rs
          Nothing -> Left $ OverlappingBounds ia a il ir
      | il >= l && ir <= r = flip fmap (go ms) $ \ms' -> bin ls l a ms' r rs
      | il <= l && ir >= r = do
          (ll, lr) <- split il ls
          (rl, rr) <- split ir rs
          pure $ bin ll il ia (bin lr l a ms r rl) ir rr
      | otherwise = Left $ LexicalError il ia ir t

insert :: Ord p => p -> a -> p -> LexTree p a -> Either (TreeError p a) (LexTree p a)
insert = insertWith (\_ _ -> Nothing)