module Data.Rope
    (
     -- * The @Rope@ type
     Rope(..),

     -- * Introducing and eliminating 'Rope's
     empty,
     singleton,

     -- * Basic interface
     append,
     null,
     length,

     -- * Reducing 'ByteString's (folds)

     -- ** Special folds
     concat,

     -- * Indexing Ropes
     index,

     depth,
     unsafeIndex,
     flatten,
     balance,
     isBalanced
) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.List (foldl1')
import Data.Monoid (Monoid, mempty, mappend, mconcat)
import Data.Word (Word8)
import Prelude hiding (concat, length, null)

leafSize = 8 -- 8 for testing, 256 for production
maxHeight = 48

-- TODO: Need invariant that Leaf is non-empty?
data Rope = Empty
          | Leaf {-# UNPACK #-} !ByteString
          | Node {-# UNPACK #-} !Int        -- length
                 {-# UNPACK #-} !Int        -- depth
                 {-# UNPACK #-} !Rope
                 {-# UNPACK #-} !Rope

instance Show Rope where
    showsPrec p = showsPrec p . flatten

instance Monoid Rope where
    mempty  = empty
    mappend = append
    mconcat = concat

-- -----------------------------------------------------------------------------
-- Introducing and eliminating 'Rope's

-- | /O(1)/ The empty 'Rope'
empty :: Rope
empty = Empty

-- | /O(1)/ Convert a 'Word8' into a 'Rope'
singleton :: Word8 -> Rope
singleton c = Leaf $ B.singleton c

-- ---------------------------------------------------------------------
-- Basic interface

-- | /O(1)/ Test whether a 'Rope' is empty.
null :: Rope -> Bool
null Empty = True
null _     = False

-- | /O(1)/ 'length' returns the length of a 'Rope' as an 'Int'.
length :: Rope -> Int
length Empty            = 0
length (Leaf ps)        = B.length ps
length (Node len _ _ _) = len

depth :: Rope -> Int
depth (Node _ depth _ _) = depth
depth _                  = 0

-- TODO: Check if this really is amortized constant time.
-- | /O(1)/ Concatenate a list of 'Rope's.
concat :: [Rope] -> Rope
concat rs = foldl append empty rs   -- FIXME? (strictness)

-- TODO: Check if this really is amortized constant time.
-- | /O(1)/ Append two 'Rope's.
append :: Rope -> Rope -> Rope
append Empty r = r
append l Empty = l
append l@(Leaf xs) r@(Leaf ys)
    | B.length xs + B.length ys <= leafSize = Leaf $ B.append xs ys
    | otherwise                             = catRopes l r
append (Leaf xs) (Node _ _ (Leaf ys) r)
    | B.length xs + B.length ys <= leafSize = catRopes (Leaf (B.append xs ys)) r
append (Node _ _ a (Leaf b)) (Leaf c)
    | B.length b + B.length c <= leafSize = catRopes a (Leaf (B.append b c))
append l r = appendAndBalance l r

appendAndBalance :: Rope -> Rope -> Rope
appendAndBalance l r
    | depth result < maxHeight = result
    | otherwise                = balance result
  where
    result = catRopes l r

type Forest = [Rope]

balance :: Rope -> Rope
balance Empty = Empty
balance r = foldl1' (flip catRopes) $ foldLeafsWith insert emptyForest r
  where
    emptyForest = []
    foldLeafsWith :: (Rope -> Forest -> Forest) -> Forest -> Rope -> Forest
    foldLeafsWith func f (Node _ _ l r) = foldLeafsWith func (foldLeafsWith func f l) r
    foldLeafsWith func f leaf = func leaf f
    insert :: Rope -> Forest -> Forest
    insert Empty fs = fs
    insert rp fs    = insert' rp fs (drop 2 fibs)

    insert' :: Rope -> Forest -> [Int] -> Forest
    insert' r [] (min:mins)
        | length r >= min = Empty : (insert' r [] mins)
        | otherwise       = [r]
    insert' r (Empty:fs) (min:mins)
        | length r >= min = Empty : (insert' r fs mins)
        | otherwise       = r : fs
    insert' r (f:fs) (min:mins) = Empty : (insert' (catRopes f r) fs mins)

catRopes :: Rope -> Rope -> Rope
catRopes Empty rp = rp
catRopes rp Empty = rp
catRopes n1@(Node len1 depth1 l1 r1) n2@(Node len2 depth2 l2 r2) =
    Node (len1 + len2) (1 + max depth1 depth2) n1 n2

isBalanced :: Rope -> Bool
isBalanced (Node len depth _ _) = len >= head (drop (depth + 1) fibs)
isBalanced _                    = True

-- ---------------------------------------------------------------------
-- Indexing ByteStrings

-- | /O(log n)/ 'Rope' index (subscript) operator, starting from 0.
index :: Rope -> Int -> Word8
index r n
    | n < 0         = moduleError "index" ("negative index: " ++ show n)
    | n >= length r = moduleError "index" ("index too large: " ++ show n
                                           ++ ", length = " ++ show (length r))
    | otherwise     = unsafeIndex r n

unsafeIndex :: Rope -> Int -> Word8
unsafeIndex Empty n = moduleError "index" ("index too large: " ++ show n
                                           ++ ", length = 0")
unsafeIndex (Leaf ps) i = B.index ps i
unsafeIndex (Node _ _ l r) n
    | n < len   = unsafeIndex l n
    | otherwise = unsafeIndex r (n - len)
  where
    len = length l

fibs = 1:1:(zipWith (+) fibs (tail fibs))

flatten :: Rope -> ByteString
flatten Empty          = B.empty
flatten (Leaf ps)      = ps
flatten (Node _ _ l r) = B.concat [flatten l, flatten r]

moduleError :: String -> String -> a
moduleError fun msg = error ("Data.Rope." ++ fun ++ ':':' ':msg)
{-# NOINLINE moduleError #-}