-- |
-- Module      : Iterate
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- @'offsetsFrom' x y@ is similar to @[x .. y - 1]@ but does not verify that
-- @x <= y@.  This removes an unnecessary branch.
--
{-# LANGUAGE MagicHash #-}
module Iterate
  ( offsets, offsetsFrom
  )
where

import Control.Exception (assert)

import GHC.Exts

{-# INLINE offsets #-}
offsets :: Int -> [Int]
offsets :: Int -> [Int]
offsets = Int -> Int -> [Int]
offsetsFrom Int
0

{-# INLINE offsetsFrom #-}
offsetsFrom :: Int -> Int -> [Int]
offsetsFrom :: Int -> Int -> [Int]
offsetsFrom xx :: Int
xx@(I# Int#
x) yy :: Int
yy@(I# Int#
y) = Bool -> [Int] -> [Int]
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
xx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yy) (Int# -> Int# -> [Int]
indices Int#
x Int#
y)

-- Implementation is derived from enumFromTo, this acts as good producer for
-- list fusion.

{-# RULES
"indices"        [~1] forall x y. indices x y = build (\c n -> indicesFB c n x y)
"indicesList"    [1] indicesFB (:) [] = indices
 #-}

indices :: Int# -> Int# -> [Int]
indices :: Int# -> Int# -> [Int]
indices Int#
x0 Int#
y = Int# -> [Int]
go Int#
x0
  where
    go :: Int# -> [Int]
go Int#
x = if Int# -> Bool
isTrue# (Int#
x Int# -> Int# -> Int#
<# Int#
y) then Int# -> Int
I# Int#
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int# -> [Int]
go (Int#
x Int# -> Int# -> Int#
+# Int#
1#) else []
{-# NOINLINE [1] indices #-}

indicesFB :: (Int -> r -> r) -> r -> Int# -> Int# -> r
indicesFB :: forall r. (Int -> r -> r) -> r -> Int# -> Int# -> r
indicesFB Int -> r -> r
c r
n Int#
x0 Int#
y = Int# -> r
go Int#
x0
  where
    go :: Int# -> r
go Int#
x = if Int# -> Bool
isTrue# (Int#
x Int# -> Int# -> Int#
<# Int#
y) then Int# -> Int
I# Int#
x Int -> r -> r
`c` Int# -> r
go (Int#
x Int# -> Int# -> Int#
+# Int#
1#) else r
n
    -- Be very careful not to have more than one "c" so that when indicesFB is
    -- inlined we can inline whatever is bound to "c"
{-# INLINE [0] indicesFB #-}