array.hs

{-# language BangPatterns #-}
{-# language ScopedTypeVariables #-}
{-# language UnboxedTuples #-}

module Array (mapMaybe) where

import Control.Monad.ST
import Data.Primitive.Array

data Alls = AllNothing | NotAllNothing

emptyMutableArray :: ST s (MutableArray s a)
emptyMutableArray = newArray 0 (error "impossible")
{-# noinline emptyMutableArray #-}

mapMaybe :: forall a b. (a -> Maybe b) -> Array a -> Array b
mapMaybe p arr = runST $ do
  let !sz = sizeofArray arr
  let go2 :: MutableArray s b -> Int -> Int -> ST s (Int, MutableArray s b)
      go2 marr !ixSrc !ixDst = if ixSrc < sz
        then do
          a <- indexArrayM arr ixSrc
          case p a of
            Just b -> do
              writeArray marr ixDst b
              go2 marr (ixSrc + 1) (ixDst + 1)
            Nothing -> go2 marr (ixSrc + 1) ixDst 
        else pure (ixDst, marr)
  let go :: Int -> Int -> Bool -> ST s (Int, MutableArray s b)
      go !ixSrc !ixDst !encountered = if ixSrc < sz
        then do
          a <- indexArrayM arr ixSrc
          case p a of
            Just b -> do
              marr <- newArray (sz - ixSrc) b
              go2 marr (ixSrc + 1) (ixDst + 1)
            Nothing -> go (ixSrc + 1) ixDst False
        else do
          marrEmpty <- emptyMutableArray
          return (ixDst, marrEmpty)
  (dstLen, marr) <- go 0 0 False  
  empty <- emptyMutableArray
  if marr == empty
    then unsafeFreezeArray marr
    -- we need resizing primitives?
    else unsafeFreezeArray marr