{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Module: Covenant.DeBruijn
-- Copyright: (C) MLabs 2025
-- License: Apache 2.0
-- Maintainer: koz@mlabs.city, sean@mlabs.city
--
-- DeBruijn indexing type and helpers. These are mainly used for scope tracking.
module Covenant.DeBruijn
  ( DeBruijn (Z, S),
    asInt,
  )
where

import Control.Monad (guard)
import Data.Coerce (coerce)
import Data.List.NonEmpty (NonEmpty)
import Data.Semigroup (Semigroup (sconcat, stimes), Sum (Sum))
import Data.Word (Word32)
import Test.QuickCheck (Arbitrary)

-- | A DeBruijn index.
--
-- @since 1.0.0
newtype DeBruijn = DeBruijn Word32
  deriving
    ( -- | @since 1.0.0
      DeBruijn -> DeBruijn -> Bool
(DeBruijn -> DeBruijn -> Bool)
-> (DeBruijn -> DeBruijn -> Bool) -> Eq DeBruijn
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DeBruijn -> DeBruijn -> Bool
== :: DeBruijn -> DeBruijn -> Bool
$c/= :: DeBruijn -> DeBruijn -> Bool
/= :: DeBruijn -> DeBruijn -> Bool
Eq,
      -- | @since 1.0.0
      Eq DeBruijn
Eq DeBruijn =>
(DeBruijn -> DeBruijn -> Ordering)
-> (DeBruijn -> DeBruijn -> Bool)
-> (DeBruijn -> DeBruijn -> Bool)
-> (DeBruijn -> DeBruijn -> Bool)
-> (DeBruijn -> DeBruijn -> Bool)
-> (DeBruijn -> DeBruijn -> DeBruijn)
-> (DeBruijn -> DeBruijn -> DeBruijn)
-> Ord DeBruijn
DeBruijn -> DeBruijn -> Bool
DeBruijn -> DeBruijn -> Ordering
DeBruijn -> DeBruijn -> DeBruijn
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DeBruijn -> DeBruijn -> Ordering
compare :: DeBruijn -> DeBruijn -> Ordering
$c< :: DeBruijn -> DeBruijn -> Bool
< :: DeBruijn -> DeBruijn -> Bool
$c<= :: DeBruijn -> DeBruijn -> Bool
<= :: DeBruijn -> DeBruijn -> Bool
$c> :: DeBruijn -> DeBruijn -> Bool
> :: DeBruijn -> DeBruijn -> Bool
$c>= :: DeBruijn -> DeBruijn -> Bool
>= :: DeBruijn -> DeBruijn -> Bool
$cmax :: DeBruijn -> DeBruijn -> DeBruijn
max :: DeBruijn -> DeBruijn -> DeBruijn
$cmin :: DeBruijn -> DeBruijn -> DeBruijn
min :: DeBruijn -> DeBruijn -> DeBruijn
Ord,
      -- | @since 1.0.0
      Gen DeBruijn
Gen DeBruijn -> (DeBruijn -> [DeBruijn]) -> Arbitrary DeBruijn
DeBruijn -> [DeBruijn]
forall a. Gen a -> (a -> [a]) -> Arbitrary a
$carbitrary :: Gen DeBruijn
arbitrary :: Gen DeBruijn
$cshrink :: DeBruijn -> [DeBruijn]
shrink :: DeBruijn -> [DeBruijn]
Arbitrary
    )
    via Word32
  deriving stock
    ( -- | @since 1.0.0
      Int -> DeBruijn -> ShowS
[DeBruijn] -> ShowS
DeBruijn -> String
(Int -> DeBruijn -> ShowS)
-> (DeBruijn -> String) -> ([DeBruijn] -> ShowS) -> Show DeBruijn
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DeBruijn -> ShowS
showsPrec :: Int -> DeBruijn -> ShowS
$cshow :: DeBruijn -> String
show :: DeBruijn -> String
$cshowList :: [DeBruijn] -> ShowS
showList :: [DeBruijn] -> ShowS
Show
    )

-- | Enables some manner of arithmetic with 'DeBruijn's. In this case, '<>' is
-- analogous to '+', while @'stimes' b@ is analogous to scalar multiplication by
-- @b@. Note that 'DeBruijn's cannot be scaled by negative numbers.
--
-- @since 1.0.0
instance Semigroup DeBruijn where
  {-# INLINEABLE (<>) #-}
  DeBruijn Word32
x <> :: DeBruijn -> DeBruijn -> DeBruijn
<> DeBruijn Word32
y = Word32 -> DeBruijn
DeBruijn (Word32
x Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
y)
  {-# INLINEABLE sconcat #-}
  sconcat :: NonEmpty DeBruijn -> DeBruijn
sconcat = Word32 -> DeBruijn
DeBruijn (Word32 -> DeBruijn)
-> (NonEmpty DeBruijn -> Word32) -> NonEmpty DeBruijn -> DeBruijn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty Word32 -> Word32
forall a. Num a => NonEmpty a -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum (NonEmpty Word32 -> Word32)
-> (NonEmpty DeBruijn -> NonEmpty Word32)
-> NonEmpty DeBruijn
-> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Coercible a b => a -> b
forall a b. Coercible a b => a -> b
coerce @(NonEmpty DeBruijn) @(NonEmpty Word32)
  {-# INLINEABLE stimes #-}
  stimes :: forall b. Integral b => b -> DeBruijn -> DeBruijn
stimes b
b = Word32 -> DeBruijn
DeBruijn (Word32 -> DeBruijn)
-> (DeBruijn -> Word32) -> DeBruijn -> DeBruijn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sum Word32 -> Word32
forall a b. Coercible a b => a -> b
coerce (Sum Word32 -> Word32)
-> (DeBruijn -> Sum Word32) -> DeBruijn -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Sum Word32 -> Sum Word32
forall b. Integral b => b -> Sum Word32 -> Sum Word32
forall a b. (Semigroup a, Integral b) => b -> a -> a
stimes b
b (Sum Word32 -> Sum Word32)
-> (DeBruijn -> Sum Word32) -> DeBruijn -> Sum Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Coercible a b => a -> b
forall a b. Coercible a b => a -> b
coerce @_ @(Sum Word32)

-- | @since 1.0.0
instance Monoid DeBruijn where
  {-# INLINEABLE mempty #-}
  mempty :: DeBruijn
mempty = DeBruijn
Z

-- | The zero index.
--
-- @since 1.0.0
pattern Z :: DeBruijn
pattern $bZ :: DeBruijn
$mZ :: forall {r}. DeBruijn -> ((# #) -> r) -> ((# #) -> r) -> r
Z <- DeBruijn 0
  where
    Z = Word32 -> DeBruijn
DeBruijn Word32
0

-- | Successor to an index.
--
-- @since 1.0.0
pattern S :: DeBruijn -> DeBruijn
pattern $bS :: DeBruijn -> DeBruijn
$mS :: forall {r}. DeBruijn -> (DeBruijn -> r) -> ((# #) -> r) -> r
S x <- (reduce -> Just x)
  where
    S (DeBruijn Word32
x) = Word32 -> DeBruijn
DeBruijn (Word32
x Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1)

{-# COMPLETE Z, S #-}

-- | Convert a DeBruijn index into a (non-negative) 'Int'.
--
-- @since 1.0.0
asInt :: DeBruijn -> Int
asInt :: DeBruijn -> Int
asInt (DeBruijn Word32
i) = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i

-- Helpers

reduce :: DeBruijn -> Maybe DeBruijn
reduce :: DeBruijn -> Maybe DeBruijn
reduce (DeBruijn Word32
x) = Word32 -> DeBruijn
DeBruijn (Word32
x Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1) DeBruijn -> Maybe () -> Maybe DeBruijn
forall a b. a -> Maybe b -> Maybe a
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ Bool -> Maybe ()
forall (f :: Type -> Type). Alternative f => Bool -> f ()
guard (Word32
x Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
> Word32
0)