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

import qualified LLVM.Extra.Multi.Value as MultiValue
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 (MultiValue.T Bool)) ->
   Iter.T r a -> Iter.T r a
takeWhile p = Iter.takeWhile (fmap unpackBool . p)

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

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

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


class (MultiValue.C a) => Enum a where
   succ, pred :: MultiValue.T a -> LLVM.CodeGenFunction r (MultiValue.T a)
   enumFrom :: MultiValue.T a -> Iter.T r (MultiValue.T a)
   enumFromTo :: MultiValue.T a -> MultiValue.T a -> Iter.T r (MultiValue.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 = MultiValue.succ
   pred = MultiValue.pred
   enumFrom = Iter.iterate MultiValue.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) =>
   MultiValue.T (Enum.T w e) ->
   MultiValue.T (Enum.T w e) ->
   LLVM.CodeGenFunction r (Maybe.T (MultiValue.T (Enum.T w e)))
succMax to e = do
   MultiValue.Cons less <- MultiValue.cmpEnum A.CmpLT e to
   C.ifThen less (Maybe.nothing Tuple.undef) $
      fmap Maybe.just $ MultiValue.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) =>
   MultiValue.T (Enum.T w e) ->
   MultiValue.T (Enum.T w e) ->
   Iter.T r (MultiValue.T (Enum.T w e))
_enumFromToSimple from to =
   takeWhile (MultiValue.cmpEnum LLVM.CmpGE to) $ enumFrom from