-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

{-# LANGUAGE FlexibleContexts, MultiParamTypeClasses, RankNTypes #-}
{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}
module Util.LazyVector (
    Lazy, lazyChunks, toStrict
    -- * builder
    , Builder, build, singleton, fromList
) where
import qualified Control.Monad.ST as ST
import qualified Control.Monad.ST.Unsafe as ST.Unsafe
import qualified Data.Vector as Vector
import qualified Data.Vector.Mutable as Mutable

import qualified Util.Pretty as Pretty


newtype Lazy a = Lazy [Vector.Vector a]
    deriving (Int -> Lazy a -> ShowS
forall a. Show a => Int -> Lazy a -> ShowS
forall a. Show a => [Lazy a] -> ShowS
forall a. Show a => Lazy a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Lazy a] -> ShowS
$cshowList :: forall a. Show a => [Lazy a] -> ShowS
show :: Lazy a -> String
$cshow :: forall a. Show a => Lazy a -> String
showsPrec :: Int -> Lazy a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Lazy a -> ShowS
Show, [Lazy a] -> Doc
Lazy a -> Text
Lazy a -> Doc
forall a. Pretty a => [Lazy a] -> Doc
forall a. Pretty a => Lazy a -> Text
forall a. Pretty a => Lazy a -> Doc
forall a. (a -> Text) -> (a -> Doc) -> ([a] -> Doc) -> Pretty a
formatList :: [Lazy a] -> Doc
$cformatList :: forall a. Pretty a => [Lazy a] -> Doc
format :: Lazy a -> Doc
$cformat :: forall a. Pretty a => Lazy a -> Doc
pretty :: Lazy a -> Text
$cpretty :: forall a. Pretty a => Lazy a -> Text
Pretty.Pretty)

instance Functor Lazy where
    fmap :: forall a b. (a -> b) -> Lazy a -> Lazy b
fmap a -> b
f (Lazy [Vector a]
chunks) = forall a. [Vector a] -> Lazy a
Lazy (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> Vector a -> Vector b
Vector.map a -> b
f) [Vector a]
chunks)

lazyChunks :: Lazy a -> [Vector.Vector a]
lazyChunks :: forall a. Lazy a -> [Vector a]
lazyChunks (Lazy [Vector a]
chunks) = [Vector a]
chunks

toStrict :: Lazy a -> Vector.Vector a
toStrict :: forall a. Lazy a -> Vector a
toStrict = forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Lazy a -> [Vector a]
lazyChunks

-- cons :: Vector.Vector a -> Lazy a -> Lazy a
-- cons v (Lazy vs)
--     | Vector.null v = Lazy vs
--     | otherwise = Lazy (v : vs)

-- * builder

newtype Builder a = Builder {
    -- Invariant: The lists include no null Vectors.
    forall a.
Builder a
-> forall s.
   (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
runBuilder :: forall s. (Buffer s a -> ST.ST s [Vector.Vector a])
         -> Buffer s a -> ST.ST s [Vector.Vector a]
    }

instance (Show a) => Show (Builder a) where
    show :: Builder a -> String
show = forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Builder a -> Lazy a
build

instance (Pretty.Pretty a) => Pretty.Pretty (Builder a) where
    format :: Builder a -> Doc
format = forall a. Pretty a => a -> Doc
Pretty.format forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Builder a -> Lazy a
build

instance Semigroup (Builder a) where
    <> :: Builder a -> Builder a -> Builder a
(<>) = forall a. Builder a -> Builder a -> Builder a
append

instance Monoid (Builder a) where
    mempty :: Builder a
mempty = forall a. Builder a
empty
    mappend :: Builder a -> Builder a -> Builder a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

build :: Builder a -> Lazy a
build :: forall a. Builder a -> Lazy a
build Builder a
builder = forall a. [Vector a] -> Lazy a
Lazy forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s a) -> a
ST.runST forall a b. (a -> b) -> a -> b
$ do
    Buffer s a
buf <- forall s a. Int -> ST s (Buffer s a)
newBuffer Int
chunkSize
    forall a.
Builder a
-> forall s.
   (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
runBuilder (Builder a
builder forall a. Builder a -> Builder a -> Builder a
`append` forall a. Builder a
flush) (forall a b. a -> b -> a
const (forall (m :: * -> *) a. Monad m => a -> m a
return [])) Buffer s a
buf

empty :: Builder a
empty :: forall a. Builder a
empty = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder (\Buffer s a -> ST s [Vector a]
cont Buffer s a
buf -> Buffer s a -> ST s [Vector a]
cont Buffer s a
buf)

append :: Builder a -> Builder a -> Builder a
append :: forall a. Builder a -> Builder a -> Builder a
append (Builder forall s.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
f) (Builder forall s.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
g) = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder (forall s.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
g)

singleton :: a -> Builder a
singleton :: forall a. a -> Builder a
singleton a
val = forall a.
Int -> (forall s. MVector s a -> Int -> ST s ()) -> Builder a
writeN Int
1 forall a b. (a -> b) -> a -> b
$ \MVector s a
mvec Int
offset ->
    forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
Mutable.unsafeWrite MVector s a
mvec Int
offset a
val

fromList :: [a] -> Builder a
fromList :: forall a. [a] -> Builder a
fromList [a]
xs = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder forall {s}.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
build
    where
    build :: (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
build Buffer s a -> ST s [Vector a]
cont (Buffer STVector s a
mvec Int
offset Int
left) = STVector s a -> Int -> Int -> [a] -> ST s [Vector a]
loop STVector s a
mvec Int
offset Int
left [a]
xs
        where
        loop :: STVector s a -> Int -> Int -> [a] -> ST s [Vector a]
loop !STVector s a
mvec !Int
offset !Int
left [] = Buffer s a -> ST s [Vector a]
cont (forall s a. STVector s a -> Int -> Int -> Buffer s a
Buffer STVector s a
mvec Int
offset Int
left)
        loop STVector s a
mvec Int
offset Int
left (a
x:[a]
xs)
            | Int
left forall a. Eq a => a -> a -> Bool
== Int
0 = do
                Vector a
chunk <- forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
Vector.unsafeFreeze STVector s a
mvec
                STVector s a
mvec <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
Mutable.new Int
chunkSize
                [Vector a]
chunks <- forall s a. ST s a -> ST s a
ST.Unsafe.unsafeInterleaveST forall a b. (a -> b) -> a -> b
$
                    STVector s a -> Int -> Int -> [a] -> ST s [Vector a]
loop STVector s a
mvec Int
0 Int
chunkSize (a
xforall a. a -> [a] -> [a]
:[a]
xs)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Vector a
chunk forall a. a -> [a] -> [a]
: [Vector a]
chunks
            | Bool
otherwise = do
                forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
Mutable.unsafeWrite STVector s a
mvec Int
offset a
x
                STVector s a -> Int -> Int -> [a] -> ST s [Vector a]
loop STVector s a
mvec (Int
offsetforall a. Num a => a -> a -> a
+Int
1) (Int
leftforall a. Num a => a -> a -> a
-Int
1) [a]
xs

flush :: Builder a
flush :: forall a. Builder a
flush = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder forall {s} {a}.
(Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
f
    where
    f :: (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
f Buffer s a -> ST s [Vector a]
cont buf :: Buffer s a
buf@(Buffer STVector s a
mvec Int
used Int
left)
        | Int
used forall a. Eq a => a -> a -> Bool
== Int
0 = Buffer s a -> ST s [Vector a]
cont Buffer s a
buf
        | Bool
otherwise = do
            Vector a
vec <- forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
Vector.unsafeFreeze STVector s a
mvec
            -- TODO why is it I can use mvec after freezing?
            -- I think because the offset writes after the frozen point.
            -- But in any case, ensureFree calls const newBuffer afterwards.
            -- Why not do the newBuffer in flush?
            let !buf2 :: Buffer s a
buf2 = forall s a. STVector s a -> Int -> Int -> Buffer s a
Buffer STVector s a
mvec Int
used Int
left
                -- The buffer may have unused space on the end, so only take
                -- the part that has data.
                !chunk :: Vector a
chunk = forall a. Int -> Vector a -> Vector a
Vector.take Int
used Vector a
vec
            -- Text uses inlineInterleaveST and says it's faster
            [Vector a]
chunks <- forall s a. ST s a -> ST s a
ST.Unsafe.unsafeInterleaveST (Buffer s a -> ST s [Vector a]
cont Buffer s a
buf2)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Vector a
chunk forall a. a -> [a] -> [a]
: [Vector a]
chunks


-- ** buffer

data Buffer s a = Buffer
    {-# UNPACK #-} !(Mutable.STVector s a)
    {-# UNPACK #-} !Int -- elements used
    {-# UNPACK #-} !Int -- space left

writeAtMost :: Int -> (forall s. Mutable.MVector s a -> Int -> ST.ST s Int)
    -> Builder a
writeAtMost :: forall a.
Int -> (forall s. MVector s a -> Int -> ST s Int) -> Builder a
writeAtMost Int
n forall s. MVector s a -> Int -> ST s Int
f = forall a. Int -> Builder a
ensureFree Int
n forall a. Builder a -> Builder a -> Builder a
`append` forall a. (forall s. Buffer s a -> ST s (Buffer s a)) -> Builder a
withBuffer (forall s a.
(MVector s a -> Int -> ST s Int) -> Buffer s a -> ST s (Buffer s a)
writeBuffer forall s. MVector s a -> Int -> ST s Int
f)

writeN :: Int -> (forall s. Mutable.MVector s a -> Int -> ST.ST s ())
    -> Builder a
writeN :: forall a.
Int -> (forall s. MVector s a -> Int -> ST s ()) -> Builder a
writeN Int
n forall s. MVector s a -> Int -> ST s ()
f = forall a.
Int -> (forall s. MVector s a -> Int -> ST s Int) -> Builder a
writeAtMost Int
n (\MVector s a
mvec Int
offset -> forall s. MVector s a -> Int -> ST s ()
f MVector s a
mvec Int
offset forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Int
n)

ensureFree :: Int -> Builder a
ensureFree :: forall a. Int -> Builder a
ensureFree !Int
n = forall a. (Int -> Builder a) -> Builder a
withSize forall a b. (a -> b) -> a -> b
$ \Int
left ->
    if Int
n forall a. Ord a => a -> a -> Bool
<= Int
left then forall a. Builder a
empty
        else forall a. Builder a
flush forall a. Builder a -> Builder a -> Builder a
`append` forall a. (forall s. Buffer s a -> ST s (Buffer s a)) -> Builder a
withBuffer (forall a b. a -> b -> a
const (forall s a. Int -> ST s (Buffer s a)
newBuffer (forall a. Ord a => a -> a -> a
max Int
n Int
chunkSize)))

writeBuffer :: (Mutable.MVector s a -> Int -> ST.ST s Int) -> Buffer s a
    -> ST.ST s (Buffer s a)
writeBuffer :: forall s a.
(MVector s a -> Int -> ST s Int) -> Buffer s a -> ST s (Buffer s a)
writeBuffer MVector s a -> Int -> ST s Int
f (Buffer MVector s a
mvec Int
used Int
left) = do
    Int
written <- MVector s a -> Int -> ST s Int
f MVector s a
mvec Int
used
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall s a. STVector s a -> Int -> Int -> Buffer s a
Buffer MVector s a
mvec (Int
used forall a. Num a => a -> a -> a
+ Int
written) (Int
left forall a. Num a => a -> a -> a
- Int
written)

withBuffer :: (forall s. Buffer s a -> ST.ST s (Buffer s a)) -> Builder a
withBuffer :: forall a. (forall s. Buffer s a -> ST s (Buffer s a)) -> Builder a
withBuffer forall s. Buffer s a -> ST s (Buffer s a)
f = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder forall a b. (a -> b) -> a -> b
$ \Buffer s a -> ST s [Vector a]
cont Buffer s a
buf -> Buffer s a -> ST s [Vector a]
cont forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s. Buffer s a -> ST s (Buffer s a)
f Buffer s a
buf

withSize :: (Int -> Builder a) -> Builder a
withSize :: forall a. (Int -> Builder a) -> Builder a
withSize Int -> Builder a
f = forall a.
(forall s.
 (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a])
-> Builder a
Builder forall a b. (a -> b) -> a -> b
$ \Buffer s a -> ST s [Vector a]
cont buf :: Buffer s a
buf@(Buffer STVector s a
_ Int
_ Int
left) ->
    forall a.
Builder a
-> forall s.
   (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
runBuilder (Int -> Builder a
f Int
left) Buffer s a -> ST s [Vector a]
cont Buffer s a
buf

newBuffer :: Int -> ST.ST s (Buffer s a)
newBuffer :: forall s a. Int -> ST s (Buffer s a)
newBuffer Int
size = do
    STVector s a
mvec <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
Mutable.new Int
size
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall s a. STVector s a -> Int -> Int -> Buffer s a
Buffer STVector s a
mvec Int
0 Int
size

-- 8 bytes for an int, so 4k = 512 elts, 1k = 128 elts
chunkSize :: Int
chunkSize :: Int
chunkSize = Int
128 forall a. Num a => a -> a -> a
- Int
chunkOverhead

-- | Overhead of one chunk, in words.
chunkOverhead :: Int
chunkOverhead :: Int
chunkOverhead = Int
4 -- not sure what MVector is, but should be this more or less