ABC 362 / Splay tree

-- | /O(n)/ Construct a vector from a 'Bundle'.
unstream :: Vector v a => Bundle v a -> v a
{-# INLINE unstream #-}
unstream s = new (New.unstream s)

-- | Construct a vector from a monadic initialiser.
new :: Vector v a => New v a -> v a
{-# INLINE_FUSED new #-}
new m = m `seq` runST (unsafeFreeze =<< New.run m)

-- ほぼ `MVector`
data New v a = New (forall s. ST s (Mutable v s a))

unstream :: Vector v a => Bundle v a -> New v a
{-# INLINE_FUSED unstream #-}
unstream s = s `seq` New (MVector.vunstream s)

-- | Create a new mutable vector and fill it with elements from the 'Bundle'.
-- The vector will grow exponentially if the maximum size of the 'Bundle' is
-- unknown.
vunstream :: (PrimMonad m, V.Vector v a)
         => Bundle v a -> m (V.Mutable v (PrimState m) a)
-- NOTE: replace INLINE_FUSED by INLINE? (also in unstreamR)
{-# INLINE_FUSED vunstream #-}
vunstream s = vmunstream (Bundle.lift s)

vmunstreamMax :: (PrimMonad m, V.Vector v a)
              => MBundle m v a -> Int -> m (V.Mutable v (PrimState m) a)
{-# INLINE vmunstreamMax #-}
vmunstreamMax s n
  = do
      v <- checkLength Internal n $ unsafeNew n
      let {-# INLINE_INNER copyChunk #-}
          copyChunk i (Chunk m f) =
            checkSlice Internal i m (length v) $ do
              f (basicUnsafeSlice i m v)
              return (i+m)

      n' <- Stream.foldlM' copyChunk 0 (MBundle.chunks s)
      return $ checkSlice Internal 0 n' n
             $ unsafeSlice 0 n' v

vmunstreamUnknown :: (PrimMonad m, V.Vector v a)
                 => MBundle m v a -> m (V.Mutable v (PrimState m) a)
{-# INLINE vmunstreamUnknown #-}
vmunstreamUnknown s
  = do
      v <- unsafeNew 0
      (v', n) <- Stream.foldlM copyChunk (v,0) (MBundle.chunks s)
      return $ checkSlice Internal 0 n (length v')
             $ unsafeSlice 0 n v'
  where
    {-# INLINE_INNER copyChunk #-}
    copyChunk (v,i) (Chunk n f)
      = do
          let j = i+n
          v' <- if basicLength v < j
                  then unsafeGrow v (delay_inline max (enlarge_delta v) (j - basicLength v))
                  else return v
          checkSlice Internal i n (length v') $ f (basicUnsafeSlice i n v')
          return (v',j)

fromVectors :: forall m v a. (Monad m, Vector v a) => [v a] -> Bundle m v a
{-# INLINE_FUSED fromVectors #-}
fromVectors us = Bundle (Stream pstep (Left us))
                        (Stream vstep us)
                        Nothing
                        (Exact n) -- ***** これ
  where
    n = List.foldl' (\k v -> k + basicLength v) 0 us

    pstep (Left []) = return Done
    pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs)))

    pstep (Right (v,i,vs))
      | i >= basicLength v = return $ Skip (Left vs)
      | otherwise          = case basicUnsafeIndexM v i of
                               Box x -> return $ Yield x (Right (v,i+1,vs))

    -- FIXME: work around bug in GHC 7.6.1
    vstep :: HasCallStack => [v a] -> m (Step [v a] (Chunk v a))
    vstep [] = return Done
    vstep (v:vs) = return $ Yield (Chunk (basicLength v)
                                         (\mv -> check
                                                 Internal
                                                 "length mismatch"
                                                 (M.basicLength mv == basicLength v)
                                                 $ stToPrim $ basicUnsafeCopy mv v)) vs

concatVectors :: (Monad m, Vector v a) => Bundle m u (v a) -> Bundle m v a
{-# INLINE_FUSED concatVectors #-}
concatVectors Bundle{sElems = Stream step t}
  = Bundle (Stream pstep (Left t))
           (Stream vstep t)
           Nothing
           Unknown -- ***** これ
  where
    pstep (Left s) = do
      r <- step s
      case r of
        Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s')))
        Skip    s' -> return (Skip (Left s'))
        Done       -> return Done

    pstep (Right (v,i,s))
      | i >= basicLength v = return (Skip (Left s))
      | otherwise          = case basicUnsafeIndexM v i of
                               Box x -> return (Yield x (Right (v,i+1,s)))


    vstep s = do
      r <- step s
      case r of
        Yield v s' -> return (Yield (Chunk (basicLength v)
                                           (\mv -> check
                                                   Internal
                                                   "length mismatch"
                                                   (M.basicLength mv == basicLength v)
                                                   $ stToPrim $ basicUnsafeCopy mv v)) s')
        Skip    s' -> return (Skip s')
        Done       -> return Done