Iterator.hs

{-# LANGUAGE TypeFamilies #-}
module LLVM.Extra.Nice.Iterator (
   takeWhile,
   countDown,
   take,
   Enum(..),
   ) where

import qualified LLVM.Extra.Nice.Value as NiceValue
import qualified LLVM.Extra.Iterator as Iter
import qualified LLVM.Extra.ScalarOrVector as SoV
import qualified LLVM.Extra.Tuple as Tuple
import qualified LLVM.Extra.MaybePrivate as Maybe
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Control as C

import qualified LLVM.Core as LLVM
import LLVM.Core (CodeGenFunction)

import Control.Applicative (liftA2)

import qualified Data.Enum.Storable as Enum

import qualified Prelude as P
import Prelude hiding (take, takeWhile, Enum, enumFrom, enumFromTo)



takeWhile ::
   (a -> CodeGenFunction r (NiceValue.T Bool)) ->
   Iter.T r a -> Iter.T r a
takeWhile p = Iter.takeWhile (fmap unpackBool . p)

unpackBool :: NiceValue.T Bool -> LLVM.Value Bool
unpackBool (NiceValue.Cons b) = b

countDown ::
   (NiceValue.Additive i, NiceValue.Comparison i,
    NiceValue.IntegerConstant i) =>
   NiceValue.T i -> Iter.T r (NiceValue.T i)
countDown len =
   takeWhile (NiceValue.cmp LLVM.CmpLT NiceValue.zero) $
   Iter.iterate NiceValue.dec len

take ::
   (NiceValue.Additive i, NiceValue.Comparison i,
    NiceValue.IntegerConstant i) =>
   NiceValue.T i -> Iter.T r a -> Iter.T r a
take len xs = liftA2 const xs (countDown len)


class (NiceValue.C a) => Enum a where
   succ, pred :: NiceValue.T a -> LLVM.CodeGenFunction r (NiceValue.T a)
   enumFrom :: NiceValue.T a -> Iter.T r (NiceValue.T a)
   enumFromTo :: NiceValue.T a -> NiceValue.T a -> Iter.T r (NiceValue.T a)

instance
   (LLVM.IsInteger w, SoV.IntegerConstant w, Num w,
    LLVM.CmpRet w, LLVM.IsPrimitive w, P.Enum e) =>
      Enum (Enum.T w e) where
   succ = NiceValue.succ
   pred = NiceValue.pred
   enumFrom = Iter.iterate NiceValue.succ
   {- |
   More complicated than 'enumFromToSimple'
   but works also for e.g. [0 .. (0xFFFF::Word16)].
   -}
   enumFromTo from to =
      Iter.takeWhileJust $
      Iter.iterate (Maybe.maybeArg Tuple.undef (succMax to)) (Maybe.just from)

succMax ::
   (LLVM.IsInteger w, SoV.IntegerConstant w, Num w,
    LLVM.CmpRet w, LLVM.IsPrimitive w, P.Enum e) =>
   NiceValue.T (Enum.T w e) ->
   NiceValue.T (Enum.T w e) ->
   LLVM.CodeGenFunction r (Maybe.T (NiceValue.T (Enum.T w e)))
succMax to e = do
   NiceValue.Cons less <- NiceValue.cmpEnum A.CmpLT e to
   C.ifThen less (Maybe.nothing Tuple.undef) $
      fmap Maybe.just $ NiceValue.succ e

{- |
Warning: For [0 .. (0xFFFF::Word16)]
it would compute an undefined @0xFFFF+1@.
In modulo arithmetic it would enter an infinite loop.
-}
_enumFromToSimple ::
   (LLVM.IsInteger w, SoV.IntegerConstant w, Num w,
    LLVM.CmpRet w, LLVM.IsPrimitive w, P.Enum e) =>
   NiceValue.T (Enum.T w e) ->
   NiceValue.T (Enum.T w e) ->
   Iter.T r (NiceValue.T (Enum.T w e))
_enumFromToSimple from to =
   takeWhile (NiceValue.cmpEnum LLVM.CmpGE to) $ enumFrom from