natural-arithmetic-0.2.0.0: src/Arithmetic/Fin.hs

{-# language BangPatterns #-}
{-# language DataKinds #-}
{-# language ExplicitNamespaces #-}
{-# language GADTs #-}
{-# language KindSignatures #-}
{-# language MagicHash #-}
{-# language PatternSynonyms #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language TypeApplications #-}
{-# language TypeOperators #-}
{-# language UnboxedTuples #-}
module Arithmetic.Fin
  ( -- * Modification
    incrementL
  , incrementR
  , incrementR#
  , weaken
  , weakenL
  , weakenR
  , succ
  , succ#
    -- * Traverse
    -- | These use the terms @ascend@ and @descend@ rather than the
    -- more popular @l@ (left) and @r@ (right) that pervade the Haskell
    -- ecosystem. The general rule is that ascending functions pair
    -- the initial accumulator with zero with descending functions
    -- pair the initial accumulator with the last index.
  , ascend
  , ascend'
  , ascendFrom'
  , ascendFrom'#
  , ascendM
  , ascendM#
  , ascendM_
  , ascendM_#
  , descend
  , descend#
  , descend'
  , descendM
  , descendM_
  , ascending
  , descending
  , ascendingSlice
  , descendingSlice
    -- * Absurdities
  , absurd
    -- * Demote
  , demote
  , demote#
    -- * Deconstruct
  , with
  , with#
    -- * Construct
  , construct#
  , remInt#
  , remWord#
    -- * Lift and Unlift
  , lift
  , unlift
  ) where

import Prelude hiding (last,succ)

import Arithmetic.Nat ((<?))
import Arithmetic.Types (Fin(..),Difference(..),Nat,Nat#,type (<), type (<#), type (<=), type (:=:))
import Arithmetic.Types (pattern MaybeFinNothing#, pattern MaybeFinJust#)
import Arithmetic.Unsafe (Nat#(Nat#),Fin#(Fin#),MaybeFin#)
import GHC.Exts (Int(I#),(+#),Int#,Word#)
import GHC.TypeNats (type (+))

import qualified Arithmetic.Lt as Lt
import qualified Arithmetic.Lte as Lte
import qualified Arithmetic.Equal as Eq
import qualified Arithmetic.Nat as Nat
import qualified Arithmetic.Plus as Plus
import qualified Arithmetic.Unsafe as Unsafe
import qualified GHC.Exts as Exts

-- | Raise the index by @m@ and weaken the bound by @m@, adding
-- @m@ to the right-hand side of @n@.
incrementR :: forall n m. Nat m -> Fin n -> Fin (n + m)
{-# inline incrementR #-}
incrementR m (Fin i pf) = Fin (Nat.plus i m) (Lt.incrementR @m pf)

incrementR# :: forall n m. Nat# m -> Fin# n -> Fin# (n + m)
{-# inline incrementR# #-}
incrementR# (Nat# n) (Fin# i) = Fin# (n +# i)

-- | Raise the index by @m@ and weaken the bound by @m@, adding
-- @m@ to the left-hand side of @n@.
incrementL :: forall n m. Nat m -> Fin n -> Fin (m + n)
{-# inline incrementL #-}
incrementL m (Fin i pf) = Fin (Nat.plus m i) (Lt.incrementL @m pf)

-- | Weaken the bound by @m@, adding it to the left-hand side of
-- the existing bound. This does not change the index.
weakenL :: forall n m. Fin n -> Fin (m + n)
{-# inline weakenL #-}
weakenL (Fin i pf) = Fin i
  ( Lt.substituteR
    (Plus.commutative @n @m)
    (Lt.plus pf (Lte.zero @m))
  )

-- | Weaken the bound by @m@, adding it to the right-hand side of
-- the existing bound. This does not change the index.
weakenR :: forall n m. Fin n -> Fin (n + m)
{-# inline weakenR #-}
weakenR (Fin i pf) = Fin i (Lt.plus pf Lte.zero)

-- | Weaken the bound, replacing it by another number greater than
-- or equal to itself. This does not change the index.
weaken :: forall n m. (n <= m) -> Fin n -> Fin m
{-# inline weaken #-}
weaken lt (Fin i pf) = Fin i (Lt.transitiveNonstrictR pf lt)

-- | A finite set of no values is impossible.
absurd :: Fin 0 -> void
{-# inline absurd #-}
absurd (Fin _ pf) = Lt.absurd pf

-- | Fold over the numbers bounded by @n@ in descending
-- order. This is lazy in the accumulator. For convenince,
-- this differs from @foldr@ in the order of the parameters.
--
-- > descend 4 z f = f 0 (f 1 (f 2 (f 3 z)))
descend :: forall a n.
     Nat n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin n -> a -> a) -- ^ Update accumulator
  -> a
{-# inline descend #-}
descend !n b0 f = go Nat.zero
  where
  go :: Nat m -> a
  go !m = case m <? n of
    Nothing -> b0
    Just lt -> f (Fin m lt) (go (Nat.succ m))

descend# :: forall a n.
     Nat# n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin# n -> a -> a) -- ^ Update accumulator
  -> a
{-# inline descend# #-}
descend# !n b0 f = descend (Nat.lift n) b0 (\ix a -> f (unlift ix) a)

-- | Fold over the numbers bounded by @n@ in descending
-- order. This is strict in the accumulator. For convenince,
-- this differs from @foldr'@ in the order of the parameters.
--
-- > descend 4 z f = f 0 (f 1 (f 2 (f 3 z)))
descend' :: forall a n.
     Nat n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin n -> a -> a) -- ^ Update accumulator
  -> a
{-# inline descend' #-}
descend' !n !b0 f = go n Lte.reflexive b0
  where
    go :: Nat p -> p <= n -> a -> a
    go !m pLteEn !b = case Nat.monus m Nat.one of
      Nothing -> b
      Just (Difference (mpred :: Nat c) cPlusOneEqP) ->
        let !cLtEn = descendLemma cPlusOneEqP pLteEn
        in go mpred (Lte.fromStrict cLtEn) (f (Fin mpred cLtEn) b)

-- | Fold over the numbers bounded by @n@ in ascending order. This
-- is lazy in the accumulator.
--
-- > ascend 4 z f = f 3 (f 2 (f 1 (f 0 z)))
ascend :: forall a n.
     Nat n
  -> a
  -> (Fin n -> a -> a)
  -> a
{-# inline ascend #-}
ascend !n !b0 f = go n Lte.reflexive
  where
    go :: Nat p -> (p <= n) -> a
    go !m pLteEn = case Nat.monus m Nat.one of
      Nothing -> b0
      Just (Difference (mpred :: Nat c) cPlusOneEqP) ->
        let !cLtEn = descendLemma cPlusOneEqP pLteEn
        in f (Fin mpred cLtEn) (go mpred (Lte.fromStrict cLtEn))

-- | Strict fold over the numbers bounded by @n@ in ascending
-- order. For convenince, this differs from @foldl'@ in the
-- order of the parameters.
--
-- > ascend' 4 z f = f 3 (f 2 (f 1 (f 0 z)))
ascend' :: forall a n.
     Nat n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin n -> a -> a) -- ^ Update accumulator
  -> a
{-# inline ascend' #-}
ascend' !n !b0 f = go Nat.zero b0
  where
  go :: Nat m -> a -> a
  go !m !b = case m <? n of
    Nothing -> b
    Just lt -> go (Nat.succ m) (f (Fin m lt) b)

-- | Generalization of @ascend'@ that lets the caller pick the starting index:
--
-- > ascend' === ascendFrom' 0
ascendFrom' :: forall a m n.
     Nat m -- ^ Index to start at
  -> Nat n -- ^ Number of steps to take
  -> a -- ^ Initial accumulator
  -> (Fin (m + n) -> a -> a) -- ^ Update accumulator
  -> a
{-# inline ascendFrom' #-}
ascendFrom' !m0 !n !b0 f = go m0 b0
  where
  end = Nat.plus m0 n
  go :: Nat k -> a -> a
  go !m !b = case m <? end of
    Nothing -> b
    Just lt -> go (Nat.succ m) (f (Fin m lt) b)

-- | Variant of @ascendFrom'@ with unboxed arguments.
ascendFrom'# :: forall a m n.
     Nat# m -- ^ Index to start at
  -> Nat# n -- ^ Number of steps to take
  -> a -- ^ Initial accumulator
  -> (Fin# (m + n) -> a -> a) -- ^ Update accumulator
  -> a
{-# inline ascendFrom'# #-}
ascendFrom'# !m0 !n !b0 f = ascendFrom' (Nat.lift m0) (Nat.lift n) b0 (\ix -> f (unlift ix))

-- | Strict monadic left fold over the numbers bounded by @n@
-- in ascending order. Roughly:
--
-- > ascendM 4 z0 f =
-- >   f 0 z0 >>= \z1 ->
-- >   f 1 z1 >>= \z2 ->
-- >   f 2 z2 >>= \z3 ->
-- >   f 3 z3
ascendM :: forall m a n. Monad m
  => Nat n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin n -> a -> m a) -- ^ Update accumulator
  -> m a
{-# inline ascendM #-}
ascendM !n !b0 f = go Nat.zero b0
  where
  go :: Nat p -> a -> m a
  go !m !b = case m <? n of
    Nothing -> pure b
    Just lt -> go (Nat.succ m) =<< f (Fin m lt) b

-- | Variant of @ascendM@ that takes an unboxed Nat and provides
-- an unboxed Fin to the callback.
ascendM# :: forall m a n. Monad m
  => Nat# n -- ^ Upper bound
  -> a -- ^ Initial accumulator
  -> (Fin# n -> a -> m a) -- ^ Update accumulator
  -> m a
{-# inline ascendM# #-}
ascendM# n !a0 f = ascendM (Nat.lift n) a0 (\ix a -> f (unlift ix) a)

-- | Monadic traversal of the numbers bounded by @n@
-- in ascending order.
--
-- > ascendM_ 4 f = f 0 *> f 1 *> f 2 *> f 3
ascendM_ :: forall m a n. Applicative m
  => Nat n -- ^ Upper bound
  -> (Fin n -> m a) -- ^ Effectful interpretion
  -> m ()
{-# inline ascendM_ #-}
ascendM_ !n f = go Nat.zero
  where
  go :: Nat p -> m ()
  go !m = case m <? n of
    Nothing -> pure ()
    Just lt -> f (Fin m lt) *> go (Nat.succ m)

-- | Variant of @ascendM_@ that takes an unboxed Nat and provides
-- an unboxed Fin to the callback.
ascendM_# :: forall m a n. Monad m
  => Nat# n -- ^ Upper bound
  -> (Fin# n -> m a) -- ^ Update accumulator
  -> m ()
{-# inline ascendM_# #-}
ascendM_# n f = ascendM_ (Nat.lift n) (\ix -> f (unlift ix))

descendLemma :: forall a b c. a + 1 :=: b -> b <= c -> a < c
{-# inline descendLemma #-}
descendLemma !aPlusOneEqB !bLteC = id
  $ Lt.transitiveNonstrictR
      (Lt.substituteR (Plus.commutative @1 @a)
      (Lt.plus Lt.zero Lte.reflexive))
  $ Lte.substituteL (Eq.symmetric aPlusOneEqB) bLteC

-- | Strict monadic left fold over the numbers bounded by @n@
-- in descending order. Roughly:
--
-- > descendM 4 z f =
-- >   f 3 z0 >>= \z1 ->
-- >   f 2 z1 >>= \z2 ->
-- >   f 1 z2 >>= \z3 ->
-- >   f 0 z3
descendM :: forall m a n. Monad m
  => Nat n
  -> a
  -> (Fin n -> a -> m a)
  -> m a
{-# inline descendM #-}
descendM !n !b0 f = go n Lte.reflexive b0
  where
    go :: Nat p -> p <= n -> a -> m a
    go !m pLteEn !b = case Nat.monus m Nat.one of
      Nothing -> pure b
      Just (Difference (mpred :: Nat c) cPlusOneEqP) ->
        let !cLtEn = descendLemma cPlusOneEqP pLteEn
        in go mpred (Lte.fromStrict cLtEn) =<< f (Fin mpred cLtEn) b

-- | Monadic traversal of the numbers bounded by @n@
-- in descending order.
--
-- > descendM_ 4 f = f 3 *> f 2 *> f 1 *> f 0
descendM_ :: forall m a n. Applicative m
  => Nat n -- ^ Upper bound
  -> (Fin n -> m a) -- ^ Effectful interpretion
  -> m ()
{-# inline descendM_ #-}
descendM_ !n f = go n Lte.reflexive
  where
  go :: Nat p -> p <= n -> m ()
  go !m !pLteEn = case Nat.monus m Nat.one of
    Nothing -> pure ()
    Just (Difference (mpred :: Nat c) cPlusOneEqP) ->
      let !cLtEn = descendLemma cPlusOneEqP pLteEn
      in f (Fin mpred cLtEn) *> go mpred (Lte.fromStrict cLtEn)

-- | Generate all values of a finite set in ascending order.
--
-- >>> ascending (Nat.constant @3)
-- [Fin 0,Fin 1,Fin 2]
ascending :: forall n. Nat n -> [Fin n]
ascending !n = go Nat.zero
  where
  go :: Nat m -> [Fin n]
  go !m = case m <? n of
    Nothing -> []
    Just lt -> Fin m lt : go (Nat.succ m)

-- | Generate all values of a finite set in descending order.
--
-- >>> descending (Nat.constant @3)
-- [Fin 2,Fin 1,Fin 0]
descending :: forall n. Nat n -> [Fin n]
descending !n = go n Lte.reflexive
  where
    go :: Nat p -> (p <= n) -> [Fin n]
    go !m !pLteEn = case Nat.monus m Nat.one of
      Nothing -> []
      Just (Difference (mpred :: Nat c) cPlusOneEqP) ->
        let !cLtEn = descendLemma cPlusOneEqP pLteEn
        in Fin mpred cLtEn : go mpred (Lte.fromStrict cLtEn)

-- | Generate 'len' values starting from 'off' in ascending order.
--
-- >>> ascendingSlice (Nat.constant @2) (Nat.constant @3) (Lte.constant @_ @6)
-- [Fin 2,Fin 3,Fin 4]
ascendingSlice
  :: forall n off len
  .  Nat off
  -> Nat len
  -> off + len <= n
  -> [Fin n]
{-# inline ascendingSlice #-}
ascendingSlice off len !offPlusLenLteEn = go Nat.zero
  where
    go :: Nat m -> [Fin n]
    go !m = case m <? len of
      Nothing -> []
      Just emLtLen ->
        let !offPlusEmLtOffPlusLen = Lt.incrementL @off emLtLen
            !offPlusEmLtEn = Lt.transitiveNonstrictR offPlusEmLtOffPlusLen offPlusLenLteEn
         in Fin (Nat.plus off m) offPlusEmLtEn : go (Nat.succ m)

-- | Generate 'len' values starting from 'off + len - 1' in descending order.
--
-- >>> descendingSlice (Nat.constant @2) (Nat.constant @3) (Lt.constant @6)
-- [Fin 4,Fin 3,Fin 2]
descendingSlice
  :: forall n off len
  .  Nat off
  -> Nat len
  -> off + len <= n
  -> [Fin n]
{-# inline descendingSlice #-}
descendingSlice !off !len !offPlusLenLteEn =
  go len Lte.reflexive
  where
    go :: Nat m -> m <= len -> [Fin n]
    go !m !mLteEn = case Nat.monus m Nat.one of
      Nothing -> []
      Just (Difference (mpred :: Nat c) cPlusOneEqEm) ->
        let !cLtLen = Lt.transitiveNonstrictR
              (Lt.substituteR (Plus.commutative @1 @c) (Lt.plus Lt.zero Lte.reflexive))
              -- c < c + 1
              (Lte.substituteL (Eq.symmetric cPlusOneEqEm) mLteEn)
              -- c + 1 <= len
            !cPlusOffLtEn = Lt.transitiveNonstrictR
              (Lt.substituteR
                (Plus.commutative @len @off)
                (Lt.plus cLtLen (Lte.reflexive @off)))
              -- c + off < off + len
              offPlusLenLteEn
        in Fin (mpred `Nat.plus` off) cPlusOffLtEn : go mpred (Lte.fromStrict cLtLen)

-- | Extract the 'Int' from a 'Fin n'. This is intended to be used
-- at a boundary where a safe interface meets the unsafe primitives
-- on top of which it is built.
demote :: Fin n -> Int
{-# inline demote #-}
demote (Fin i _) = Nat.demote i

demote# :: Fin# n -> Int#
{-# inline demote# #-}
demote# (Fin# i) = i

lift :: Unsafe.Fin# n -> Fin n
{-# inline lift #-}
lift (Unsafe.Fin# i) = Fin (Unsafe.Nat (I# i)) Unsafe.Lt

unlift :: Fin n -> Unsafe.Fin# n
{-# inline unlift #-}
unlift (Fin (Unsafe.Nat (I# i)) _) = Unsafe.Fin# i

-- | Consume the natural number and the proof in the Fin.
with :: Fin n -> (forall i. (i < n) -> Nat i -> a) -> a
{-# inline with #-}
with (Fin i lt) f = f lt i

-- | Variant of 'with' for unboxed argument and result types.
with# :: Fin# n -> (forall i. (i <# n) -> Nat# i -> a) -> a
{-# inline with# #-}
with# (Unsafe.Fin# i) f = f (Unsafe.Lt# (# #)) (Unsafe.Nat# i)

construct# :: (i <# n) -> Nat# i -> Fin# n
{-# inline construct# #-}
construct# _ (Unsafe.Nat# x) = Unsafe.Fin# x

-- | Return the successor of the Fin or return nothing if the
-- argument is the greatest inhabitant.
succ :: Nat n -> Fin n -> Maybe (Fin n)
{-# inline succ #-}
succ n (Fin ix _) = case ix' <? n of
  Nothing -> Nothing
  Just lt -> Just (Fin ix' lt)
  where
  ix' = Nat.succ ix

-- | Variant of 'succ' for unlifted finite numbers.
succ# :: Nat# n -> Fin# n -> MaybeFin# n
{-# inline succ# #-}
succ# (Nat# n) (Fin# ix) = case ix' Exts.<# n of
  0# -> MaybeFinNothing#
  _ -> MaybeFinJust# (Fin# ix')
  where
  !ix' = ix +# 1#

-- | This crashes if @n = 0@. Divides @i@ by @n@ and takes
-- the remainder.
remInt# :: Int# -> Nat# n -> Fin# n
remInt# i (Nat# n) = case n of
  0# -> errorWithoutStackTrace "Arithmetic.Fin.remInt#: cannot divide by zero"
  _ -> Fin# (Exts.remInt# i n)

-- | This crashes if @n = 0@. Divides @i@ by @n@ and takes
-- the remainder.
remWord# :: Word# -> Nat# n -> Fin# n
remWord# w (Nat# n) = case n of
  0# -> errorWithoutStackTrace "Arithmetic.Fin.remWord#: cannot divide by zero"
  _ -> Fin# (Exts.word2Int# (Exts.remWord# w (Exts.int2Word# n)))