Index.hs

{-# LANGUAGE RoleAnnotations #-}

-- |
-- Module: Covenant.Index
-- Copyright: (C) MLabs 2025
-- License: Apache 2.0
-- Maintainer: koz@mlabs.city, sean@mlabs.city
--
-- Positional indexes, starting from 0, and cardinality indicators.
--
-- @since 1.0.0
module Covenant.Index
  ( Index,
    Count,
    intIndex,
    intCount,
    ix0,
    count0,
    ix1,
    count1,
    ix2,
    count2,
    ix3,
    count3,
    wordCount,
  )
where

import Data.Bits (toIntegralSized)
import Data.Coerce (coerce)
import Data.List.NonEmpty (NonEmpty)
import Data.Semigroup (Semigroup (sconcat, stimes), Sum (Sum))
import Data.Word (Word32)
import GHC.TypeLits (Symbol)
import Optics.Core (Lens', lens)
import Optics.Prism (Prism', prism)
import Test.QuickCheck (Arbitrary)

-- | A positional index, starting from zero. The label allows distinguishing
-- different flavours of indices.
--
-- @since 1.0.0
newtype Index (ofWhat :: Symbol) = Index Word32
  deriving
    ( -- | @since 1.0.0
      Eq,
      -- | @since 1.0.0
      Ord,
      -- | @since 1.0.0
      Arbitrary
    )
    via Word32
  deriving stock
    ( -- | @since 1.0.0
      Show
    )

type role Index nominal

-- | Enables some manner of arithmetic with 'Index'ess. In this case, '<>' is
-- analogous to '+', while @'stimes' b@ is analogous to scalar multiplication by
-- @b@. Note that 'Index'es cannot be scaled by negative numbers.
--
-- @since 1.0.0
instance Semigroup (Index ofWhat) where
  {-# INLINEABLE (<>) #-}
  Index x <> Index y = Index (x + y)
  {-# INLINEABLE sconcat #-}
  sconcat = Index . sum . coerce @(NonEmpty (Index ofWhat)) @(NonEmpty Word32)
  {-# INLINEABLE stimes #-}
  stimes b = Index . coerce . stimes b . coerce @_ @(Sum Word32)

-- | @since 1.0.0
instance Monoid (Index ofWhat) where
  {-# INLINEABLE mempty #-}
  mempty = Index 0

-- | Helper to construct, and convert, 'Index'es and 'Int's. This is needed
-- because unfortunately, the standard Haskell practice is to use 'Int' for
-- indexes.
--
-- To use this, do one of the following:
--
-- * Construct with @'preview'@: for example, @'preview' intIndex 1@.
-- * Destruct with @'review'@.
--
-- @since 1.0.0
intIndex :: forall (ofWhat :: Symbol). Prism' Int (Index ofWhat)
intIndex =
  prism
    (fromIntegral . coerce @_ @Word32)
    (\i -> maybe (Left i) (Right . Index) . toIntegralSized $ i)

-- | Helper for the first index.
--
-- @since 1.0.0
ix0 :: forall (ofWhat :: Symbol). Index ofWhat
ix0 = Index 0

-- | Helper for the second index.
--
-- @since 1.0.0
ix1 :: forall (ofWhat :: Symbol). Index ofWhat
ix1 = Index 1

-- | Helper for the third index.
--
-- @since 1.0.0
ix2 :: forall (ofWhat :: Symbol). Index ofWhat
ix2 = Index 2

-- | Helper for the fourth index.
--
-- @since 1.0.0
ix3 :: forall (ofWhat :: Symbol). Index ofWhat
ix3 = Index 3

-- | An indicator of the cardinality of something. Meant to be paired with
-- 'Index' to specify which unique something you mean.
--
-- @since 1.0.0
newtype Count (ofWhat :: Symbol) = Count Word32
  deriving
    ( -- | @since 1.0.0
      Eq,
      -- | @since 1.0.0
      Ord
    )
    via Word32
  deriving stock
    ( -- | @since 1.0.0
      Show
    )

type role Count nominal

-- | Helper to construct, and convert, 'Count's and 'Int's. This is needed
-- because unfortunately, sizes of things are usually 'Int's in Haskell.
--
-- To use this, do one of the following:
--
-- * Construct with @'preview'@: for example, @'preview' intCount 1@.
-- * Destruct with @'review'@.
--
-- @since 1.0.0
intCount :: forall (ofWhat :: Symbol). Prism' Int (Count ofWhat)
intCount =
  prism
    (fromIntegral . coerce @_ @Word32)
    (\i -> maybe (Left i) (Right . Count) . toIntegralSized $ i)

-- | We use the Word32 directly during renaming, and a Lens is more appropriate
-- than a Prism if we're working with Word32s
wordCount :: forall (ofWhat :: Symbol). Lens' (Count ofWhat) Word32
wordCount = lens (\(Count x) -> x) (\_ w -> Count w)

-- | Helper for a count of zero items.
--
-- @since 1.0.0
count0 :: forall (ofWhat :: Symbol). Count ofWhat
count0 = Count 0

-- | Helper for a count of one item.
--
-- @since 1.0.0
count1 :: forall (ofWhat :: Symbol). Count ofWhat
count1 = Count 1

-- | Helper for a count of two items.
--
-- @since 1.0.0
count2 :: forall (ofWhat :: Symbol). Count ofWhat
count2 = Count 2

-- | Helper for a count of three items.
--
-- @since 1.0.0
count3 :: forall (ofWhat :: Symbol). Count ofWhat
count3 = Count 3