mlkem-0.1.1.0: src/Iterate.hs

-- |
-- 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 = offsetsFrom 0

{-# INLINE offsetsFrom #-}
offsetsFrom :: Int -> Int -> [Int]
offsetsFrom xx@(I# x) yy@(I# y) = assert (xx <= yy) (indices x 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 x0 y = go x0
  where
    go x = if isTrue# (x <# y) then I# x : go (x +# 1#) else []
{-# NOINLINE [1] indices #-}

indicesFB :: (Int -> r -> r) -> r -> Int# -> Int# -> r
indicesFB c n x0 y = go x0
  where
    go x = if isTrue# (x <# y) then I# x `c` go (x +# 1#) else 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 #-}