{- GHC optimizer request:
 -
 - View + Pattern match is not optimized significantly, leading to significant
 - allocator churn when using this pattern.  I expect the view function to
 - run as a coroutine of the pattern match, allowing fusion to take place.
 -}
module PlzOptimize (consts, optConsts) where
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as B
import qualified Data.ByteString.Unsafe as B
import Data.Word

type ByteString = B.ByteString
type DNABlock   = ByteString    -- invariant: never empty
type DNA        = [DNABlock]    -- if i did this again this would be Data.ByteString.Lazy

data D = I|C|F|P

dnaConvElem :: Char -> D
dnaConvElem 'I' = I
dnaConvElem 'C' = C
dnaConvElem 'F' = F
dnaConvElem 'P' = P
dnaConvElem _   = error "unreachable"

dnaView :: DNA -> [D]
dnaView [] = []
dnaView (d : ds) = dnaConvElem (B.w2c $ B.unsafeHead d) : viewTail
   where
       tailD = B.unsafeTail d
       viewTail | B.null d  = dnaView ds
                | otherwise = dnaView (tailD : ds)

dnaTake :: Int -> DNA -> DNA
dnaTake x y | x `seq` y `seq` False = undefined
dnaTake 0 _  = []
dnaTake _ [] = []
dnaTake n (d : ds) | B.length d < n = d : (dnaTake (n - B.length d) ds)
                   | otherwise      = [B.unsafeTake n d]

dnaDrop :: Int -> DNA -> DNA
dnaDrop x y | x `seq` y `seq` False = undefined
dnaDrop 0 d  = d
dnaDrop _ [] = []
dnaDrop n (d : ds) | n < B.length d = B.unsafeDrop n d : ds
                   | otherwise      = dnaDrop (n - B.length d) ds

-- I expect the view function dnaView to get inlined into the pattern match here.
-- (optConsts demonstrates code I would expect the optimizer to generate for me)
consts :: DNA -> (DNA, DNA)
consts dna = let len = consts' (dnaView dna) 0 in (dnaTake len dna, dnaDrop len dna) where
    consts' :: [D] -> Int -> Int
    consts' (C:d)   n = consts' d $! (n+1)
    consts' (F:d)   n = consts' d $! (n+1)
    consts' (P:d)   n = consts' d $! (n+1)
    consts' (I:C:d) n = consts' d $! (n+2)
    consts' _       n = n

optConsts :: DNA -> (DNA, DNA)
optConsts dna = let len = optConsts' dna 0 in (dnaTake len dna, dnaDrop len dna) where
    optConsts' []       n = n
    optConsts' (d : ds) n = optConstsWorker d ds n
    
    optConstsWorker d ds n = case (B.w2c $ B.unsafeHead d) of
        'C' -> case B.unsafeTail d of
                    dt | B.null dt -> optConsts' ds $! (n+1)
                       | otherwise -> optConstsWorker dt ds $! (n+1)
        'F' -> case B.unsafeTail d of
                    dt | B.null dt -> optConsts' ds $! (n+1)
                       | otherwise -> optConstsWorker dt ds $! (n+1)
        'P' -> case B.unsafeTail d of
                    dt | B.null dt -> optConsts' ds $! (n+1)
                       | otherwise -> optConstsWorker dt ds $! (n+1)
        'I' -> case B.unsafeTail d of
                    dt | B.null dt -> case ds of
                                         [] -> n
                                         (d2 : ds2) -> case (B.w2c $ B.unsafeHead d2) of
                                                          'C' -> case B.unsafeTail d2 of
                                                                  dt2 | B.null dt2 -> optConsts' ds $! (n+2)
                                                                      | otherwise  -> optConstsWorker dt2 ds $! (n+2)
                                                          _   -> n
                       | otherwise -> case (B.w2c $ B.unsafeHead dt) of
                                         'C' -> case B.unsafeTail dt of
                                                    dt2 | B.null dt2 -> optConsts' ds $! (n+2)
                                                        | otherwise  -> optConstsWorker dt2 ds $! (n+2)
                                         _   -> n
        _   -> n