Internal.hs

{-# LANGUAGE FlexibleContexts, FlexibleInstances, TypeSynonymInstances, MultiParamTypeClasses, UndecidableInstances, TypeOperators, DeriveDataTypeable #-}
module Data.Rope.Internal
    ( Rope(..)
    -- * Construction
    , pack                  -- :: a `Reducer` Rope => a -> Roe
    , empty                 -- :: Rope
    , fromChunks            -- :: [ByteString] -> Rope
    , fromByteString        -- :: ByteString -> Rope
    , fromLazyByteString    -- :: L.ByteString -> Rope
    , fromString            -- :: String -> Rope
    , fromWords             -- :: [Word8] -> Rope
    , fromChar              -- :: Char -> Rope
    , fromWord8             -- :: Word8 -> Rope
    -- * Analysis
    , length                -- :: Rope -> Int
    , null                  -- :: Rope -> Bool
    -- * Deconstruction
    , toChunks              -- :: Rope -> [ByteString]
    , toString              -- :: Rope -> String
    , toLazyByteString      -- :: Rope -> L.ByteString
    -- * Cutting 
    , splitAt
    , take
    , drop
    -- * Unpacking
    , Unpackable(..)
    , Breakable(..)
    -- Utility 
    , w2c
    , findIndexOrEnd        -- :: (Word8 -> Bool) -> ByteString -> Int
    ) where


import Prelude hiding (head, last, length, foldl, null, length, splitAt, take, drop, break, span)
import qualified Prelude

import Control.Applicative hiding (empty)

import Data.Data (Data(..), DataType, Constr, Fixity(..), mkConstr, mkDataType, constrIndex)
import Data.Typeable (Typeable(..))

import Data.FingerTree (ViewL(..),ViewR(..),viewl,viewr,(<|),(|>), Measured(..), (><))
import qualified Data.FingerTree as F (empty, split, null, singleton)

import qualified Data.Foldable as F

import Data.Monoid

import Data.Rope.Body
import Data.Rope.Util.Reducer (Reducer, cons, snoc, unit)

import Data.Word (Word8)

import GHC.Base (unsafeChr)

import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Storable (peek)

import qualified Data.ByteString as S (null, splitAt, take, drop, length, singleton, unpack, last)
import Data.ByteString.Internal (ByteString(..), inlinePerformIO)
import qualified Data.ByteString.Unsafe as S (unsafeTail, unsafeHead)
import qualified Data.ByteString.UTF8 as U (fromString)
import qualified Data.ByteString.Lazy as L (ByteString, pack, fromChunks, toChunks, elemIndex)
import qualified Data.ByteString.Lazy.UTF8 as LU (fromString, toString)

import Codec.Binary.UTF8.Generic (UTF8Bytes)
import qualified Codec.Binary.UTF8.Generic as UTF8Bytes

-- a Buffer is a fingertree of non-empty chunks
newtype Rope = Rope { body :: Body } 
    deriving (Show, Typeable)

instance Monoid Rope where
    mempty = empty
    Rope t `mappend` Rope t' = Rope (t >< t')

instance Eq Rope where
    a == b = measure (body a) == measure (body b) 
          && toLazyByteString a == toLazyByteString b 

instance Ord Rope where
    a `compare` b = toLazyByteString a `compare` toLazyByteString b

instance Measured Offset Rope where
    measure = measure . body

pack :: Reducer a Rope => a -> Rope
pack = unit
{-# INLINE pack #-}

empty :: Rope
empty = Rope F.empty
{-# INLINE empty #-}

fromChunks :: [ByteString] -> Rope
fromChunks = foldr (\l (Rope t) -> Rope (l `cons'` t)) mempty
{-# INLINE fromChunks #-}

toChunks :: Rope -> [ByteString]
toChunks r = unchunk <$> F.toList (body r)
{-# INLINE toChunks #-}

toLazyByteString :: Rope -> L.ByteString
toLazyByteString = L.fromChunks . toChunks
{-# INLINE toLazyByteString #-}

toString :: Rope -> String
toString = unpack

length :: Rope -> Int
length = measureBody . body
{-# INLINE length #-}

null :: Rope -> Bool
null = F.null . body
{-# INLINE null #-}

fromByteString :: ByteString -> Rope
fromByteString b | S.null b = mempty 
                 | otherwise = Rope (F.singleton (Chunk b))
{-# INLINE fromByteString #-}

-- NB this requires a strict bytestring reducer, but a lazy bytestring
fromLazyByteString :: L.ByteString -> Rope
fromLazyByteString = foldr (\l (Rope t) -> Rope (Chunk l <| t)) mempty . L.toChunks
{-# INLINE fromLazyByteString #-}

-- utf8 encode chunks of the string
fromString :: String -> Rope
fromString = fromLazyByteString . LU.fromString
{-# INLINE fromString #-}

fromWords :: [Word8] -> Rope
fromWords = fromLazyByteString . L.pack
{-# INLINE fromWords #-}

fromChar :: Char -> Rope
fromChar c = Rope (F.singleton (Chunk (U.fromString [c])))
{-# INLINE fromChar #-}

fromWord8 :: Word8 -> Rope
fromWord8 b = Rope (F.singleton (Chunk (S.singleton b)))
{-# INLINE fromWord8 #-}

cons8 :: Word8 -> Rope -> Rope
cons8 a (Rope t) = case viewl t of
    Chunk c :< cs | S.length c < 16 -> Rope (Chunk (mappend b c) <| cs)
    _                               -> Rope (Chunk b <| t)
    where b = S.singleton a
{-# INLINE cons8 #-}

instance Data Rope where
    gfoldl f z r = case uncons8 r of
        Nothing -> z empty
        Just (x,xs) -> z cons8 `f` x `f` xs 
     
    gunfold k z c = case constrIndex c of
       1 -> z empty
       2 -> k (k (z cons8))
       _ -> error "gunfoldl"
    
    toConstr xs
       | null xs = emptyConstr
       | otherwise = consConstr

    dataTypeOf _ = ropeDataType

emptyConstr, consConstr :: Constr
emptyConstr = mkConstr ropeDataType "empty" [] Prefix
consConstr = mkConstr ropeDataType "`cons`" [] Infix

ropeDataType :: DataType
ropeDataType = mkDataType "Data.Rope.Internal.Rope" [emptyConstr, consConstr]

splitAt :: Int -> Rope -> (Rope,Rope)
splitAt n (Rope f)
        | n <= 0 = (mempty, Rope f)
        | n >= measureBody f = (Rope f, mempty)
        | otherwise = (Rope (x `snoc'` y'), Rope (y'' `cons'` z))
        where
            (x,yz) = F.split (> Offset n) f
            Chunk y :< z = viewl yz
            (y', y'') = S.splitAt (n - measureBody x) y

take :: Int -> Rope -> Rope
take n = fst . splitAt n
{-# INLINE take #-}

drop :: Int -> Rope -> Rope
drop n = snd . splitAt n
{-# INLINE drop #-}

class Breakable a where
    break :: (a -> Bool) -> Rope -> (Rope, Rope)
    span :: (a -> Bool) -> Rope -> (Rope, Rope)
    takeWhile :: (a -> Bool) -> Rope -> Rope
    dropWhile :: (a -> Bool) -> Rope -> Rope

    span f = break (not . f)
    takeWhile f = fst . span f
    dropWhile f = snd . span f

break8 :: (Word8 -> Bool) -> Rope -> (Rope, Rope)
break8 f r = (Rope t', Rope t'')
    where 
        (t',t'') = break' (body r)
        break' ccs = case viewl ccs of
           EmptyL -> (F.empty, F.empty)
           Chunk c :< cs -> case findIndexOrEnd f c of 
                0              -> (F.empty, ccs)
                n | n < S.length c -> (F.singleton (Chunk (S.take n c)), Chunk (S.drop n c) <| cs)
                  | otherwise      -> let (cs', cs'') = break' cs
                                      in (Chunk c <| cs', cs'')
{-# INLINE break8 #-}

instance Breakable Word8 where
    break = break8

findIndexOrEnd :: (Word8 -> Bool) -> ByteString -> Int
findIndexOrEnd k (PS x s l) = inlinePerformIO $ withForeignPtr x $ \f -> go (f `plusPtr` s) 0
  where
    go ptr n | ptr `seq` n `seq` False = undefined
             | n >= l    = return l
             | otherwise = do w <- peek ptr
                              if k w
                                then return n
                                else go (ptr `plusPtr` 1) (n+1)
{-# INLINE findIndexOrEnd #-}

uncons8 :: Rope -> Maybe (Word8, Rope)
uncons8 r = case viewl (body r) of
    Chunk c :< cs -> Just (S.unsafeHead c, Rope (S.unsafeTail c `cons'` cs))
    _ -> Nothing
{-# INLINE uncons8 #-}

unsnoc8 :: Rope -> Maybe (Rope, Word8)
unsnoc8 r = case viewr (body r) of
    cs :> Chunk c -> Just (Rope (cs `snoc'` S.unsafeTail c), S.unsafeHead c)
    _ -> Nothing
{-# INLINE unsnoc8 #-}

w2c :: Word8 -> Char
w2c = unsafeChr . fromIntegral
{-# INLINE w2c #-}

instance Reducer Char Rope where
    unit = fromChar
    cons a (Rope t) = case viewl t of
        Chunk c :< cs | S.length c < 16 -> Rope (Chunk (mappend b c) <| cs)
        _ -> Rope (Chunk b <| t)
        where b = U.fromString [a]
    snoc (Rope t) a = case viewr t of
        cs :> Chunk c | S.length c < 16 -> Rope (cs |> Chunk (mappend c b))
        _ -> Rope (t |> Chunk b)
        where b = U.fromString [a]

instance Reducer Word8 Rope where
    unit = fromWord8
    cons = cons8 
    snoc (Rope t) a = case viewr t of
        cs :> Chunk c | S.length c < 16 -> Rope (cs |> Chunk (mappend c b))
        _ -> Rope (t |> Chunk b)
        where b = S.singleton a

instance Reducer Rope Rope where
    unit = id

instance Reducer String Rope where
    unit = fromString

instance Reducer [Word8] Rope where
    unit = fromWords

instance Reducer ByteString Rope where
    unit = fromByteString

instance Reducer L.ByteString Rope where
    unit = fromLazyByteString

instance Reducer Chunk Rope where
    unit = fromByteString . unchunk 


instance UTF8Bytes Rope Int where
    bsplit = splitAt 
    bdrop = drop 
    buncons f = case viewl (body f) of
        Chunk c :< cs -> Just (S.unsafeHead c, Rope (S.unsafeTail c `cons' ` cs))
        EmptyL -> Nothing
    tail (Rope f) = case viewl f of
        Chunk c :< cs -> Rope (S.unsafeTail c `cons'`cs)
        EmptyL -> errorEmptyList "tail"
    elemIndex b = fmap fromIntegral . L.elemIndex b . L.fromChunks . map unchunk . F.toList . body
    pack = Rope . foldr (\l r -> Chunk l <| r) F.empty . L.toChunks . L.pack
    empty = Rope F.empty
    null = F.null . body


class Unpackable a where
    unpack :: Rope -> [a]

    head :: Rope -> a
    head = Prelude.head . unpack

    last :: Rope -> a

    uncons :: Rope -> Maybe (a, Rope)
    unsnoc :: Rope -> Maybe (Rope, a)

instance Unpackable Word8 where
    unpack = concatMap (S.unpack . unchunk) . F.toList . body
    head t = case viewl (body t) of
        Chunk a :< _ -> S.unsafeHead a
        EmptyL -> errorEmptyList "head"
    last t = case viewr (body t) of
        _ :> Chunk a -> S.last a
        EmptyR -> errorEmptyList "last"
    uncons = uncons8
    unsnoc = unsnoc8

instance Unpackable Char where
    unpack = LU.toString . toLazyByteString
    head = Prelude.head . unpack
    last = undefined -- TODO
    uncons r@(Rope t) = case UTF8Bytes.decode (Rope t) of 
        Nothing -> Nothing
        Just (a,n) -> Just (a, drop n r)
    unsnoc = undefined -- TODO

instance Unpackable ByteString where
    unpack = map unchunk . F.toList . body
    head r = case viewl (body r) of
        Chunk a :< _ -> a
        _ -> errorEmptyList "head"
    last r = case viewr (body r) of
        _ :> Chunk a -> a
        _ -> errorEmptyList "last" 
    uncons r = case viewl (body r) of
        Chunk a :< as -> Just (a, Rope as)
        EmptyL -> Nothing
    unsnoc r = case viewr (body r) of
        as :> Chunk a -> Just (Rope as, a)
        EmptyR -> Nothing

instance Unpackable Chunk where
    unpack = F.toList . body
    head r = case viewl (body r) of
        a :< _ -> a
        _ -> errorEmptyList "head"
    last r = case viewr (body r) of
        _ :> a -> a
        _ -> errorEmptyList "last"
    uncons r = case viewl (body r) of
        Chunk a :< as -> Just (Chunk a, Rope as)
        EmptyL -> Nothing
    unsnoc r = case viewr (body r) of
        as :> Chunk a -> Just (Rope as, Chunk a)
        EmptyR -> Nothing

errorEmptyList :: String -> a
errorEmptyList t = error $ "Data.Rope.Unpackable." ++ t ++ ": empty list"