{-# OPTIONS_GHC -Wunused-imports #-}

module Mikan.Utils.Suffix
  ( Suffix(..)
  , nextSuffix
  , suffixView
  , addSuffix
  )
  where

import Control.Monad.ST

import Data.Char
import Data.Primitive.ByteArray
import Data.Text qualified as T
import Data.Text.Short (ShortText)
import Data.Text.Short qualified as TS
import Data.Word

import Math.NumberTheory.Logarithms (integerLog10')

import Mikan.Utils.Text qualified as T
import Mikan.Utils.ShortText qualified as TS

------------------------------------------------------------------------
-- Subscript digits

{-# INLINE isSubscriptDigit #-}
-- | Is the character one of the subscripts @'₀'@-@'₉'@?
isSubscriptDigit :: Char -> Bool
isSubscriptDigit :: Char -> Bool
isSubscriptDigit Char
c = Char
'₀' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'₉'

toDigit :: Char -> Maybe Integer
toDigit :: Char -> Maybe Integer
toDigit Char
c
  | Char -> Bool
isDigit Char
c = Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer -> Maybe Integer) -> Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$! (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord Char
c) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
0x30)
  | Bool
otherwise = Maybe Integer
forall a. Maybe a
Nothing

toSubscriptDigit :: Char -> Maybe Integer
toSubscriptDigit :: Char -> Maybe Integer
toSubscriptDigit Char
c
  | Char -> Bool
isSubscriptDigit Char
c = Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer -> Maybe Integer) -> Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$! (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord Char
c) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
0x2080)
  | Bool
otherwise = Maybe Integer
forall a. Maybe a
Nothing

------------------------------------------------------------------------
-- Suffices

-- | Classification of identifier variants.

data Suffix
  = Prime Int
  -- ^ Identifier ends in 'Int' many primes.
  | Index Integer
  -- ^ Identifier ends in 'Natural' (ordinary digits).
  | Subscript Integer
  -- ^ Identifier ends in number 'Natural' (subscript digits).
  deriving Int -> Suffix -> ShowS
[Suffix] -> ShowS
Suffix -> String
(Int -> Suffix -> ShowS)
-> (Suffix -> String) -> ([Suffix] -> ShowS) -> Show Suffix
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Suffix -> ShowS
showsPrec :: Int -> Suffix -> ShowS
$cshow :: Suffix -> String
show :: Suffix -> String
$cshowList :: [Suffix] -> ShowS
showList :: [Suffix] -> ShowS
Show

-- | Increase the suffix by one.
nextSuffix :: Suffix -> Suffix
nextSuffix :: Suffix -> Suffix
nextSuffix (Prime Int
i)     = Int -> Suffix
Prime (Int -> Suffix) -> Int -> Suffix
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
nextSuffix (Index Integer
i)     = Integer -> Suffix
Index (Integer -> Suffix) -> Integer -> Suffix
forall a b. (a -> b) -> a -> b
$ Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
nextSuffix (Subscript Integer
i) = Integer -> Suffix
Subscript (Integer -> Suffix) -> Integer -> Suffix
forall a b. (a -> b) -> a -> b
$ Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1

{-# NOINLINE suffixView #-}
-- | Parse a suffix.
suffixView :: ShortText -> (ShortText, Maybe Suffix)
suffixView :: ShortText -> (ShortText, Maybe Suffix)
suffixView ShortText
st =
  let t :: Text
t = ShortText -> Text
TS.toText ShortText
st
  in case Text -> Maybe (Text, Char)
T.unsnoc Text
t of
    Just (Text
t, Char
'\'') ->
      let (Text
t', Text
primes) = (Char -> Bool) -> Text -> (Text, Text)
T.spanEnd (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\'') Text
t
          !ts :: ShortText
ts = Text -> ShortText
TS.fromText Text
t'
          !suff :: Maybe Suffix
suff = Suffix -> Maybe Suffix
forall a. a -> Maybe a
Just (Suffix -> Maybe Suffix) -> Suffix -> Maybe Suffix
forall a b. (a -> b) -> a -> b
$! Int -> Suffix
Prime (Int -> Suffix) -> Int -> Suffix
forall a b. (a -> b) -> a -> b
$! Text -> Int
T.length Text
primes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      in (ShortText
ts, Maybe Suffix
suff)
    Just (Text
t, Char -> Maybe Integer
toDigit -> Just Integer
d) -> Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t Integer
d Integer
10 where
      loop :: Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t !Integer
n !Integer
k =
        case Text -> Maybe (Text, Char)
T.unsnoc Text
t of
          Just (Text
t, Char -> Maybe Integer
toDigit -> Just Integer
d) -> Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n) (Integer
10 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k)
          Maybe (Text, Char)
_ ->
            let !ts :: ShortText
ts = Text -> ShortText
TS.fromText Text
t
                !suff :: Maybe Suffix
suff = Suffix -> Maybe Suffix
forall a. a -> Maybe a
Just (Suffix -> Maybe Suffix) -> Suffix -> Maybe Suffix
forall a b. (a -> b) -> a -> b
$! Integer -> Suffix
Index Integer
n
            in (ShortText
ts, Maybe Suffix
suff)
    Just (Text
t, Char -> Maybe Integer
toSubscriptDigit -> Just Integer
d) -> Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t Integer
d Integer
10 where
      loop :: Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t !Integer
n !Integer
k =
        case Text -> Maybe (Text, Char)
T.unsnoc Text
t of
          Just (Text
t, Char -> Maybe Integer
toSubscriptDigit -> Just Integer
d) -> Text -> Integer -> Integer -> (ShortText, Maybe Suffix)
loop Text
t (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n) (Integer
10 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k)
          Maybe (Text, Char)
_ ->
            let !ts :: ShortText
ts = Text -> ShortText
TS.fromText Text
t
                !suff :: Maybe Suffix
suff = Suffix -> Maybe Suffix
forall a. a -> Maybe a
Just (Suffix -> Maybe Suffix) -> Suffix -> Maybe Suffix
forall a b. (a -> b) -> a -> b
$! Integer -> Suffix
Subscript Integer
n
            in (ShortText
ts, Maybe Suffix
suff)
    Maybe (Text, Char)
_ -> (ShortText
st, Maybe Suffix
forall a. Maybe a
Nothing)


-- | Add a suffix onto the end of some 'ShortText'.
--
-- == __Performance__
--
-- This function is actually quite hot, as it gets invoked
-- on every name during during goal display. As such,
-- it has been optimized to avoid intermediate allocations
-- as much as possible.
addSuffix :: ShortText -> Suffix -> ShortText
addSuffix :: ShortText -> Suffix -> ShortText
addSuffix ShortText
str (Prime Int
n) = do
  let len :: Int
len = ShortText -> Int
TS.lengthBytes ShortText
str
  Int -> (forall s. MutableByteArray s -> ST s ()) -> ShortText
TS.unsafeCreate (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) \MutableByteArray s
arr -> do
    MutableByteArray (PrimState (ST s))
-> Int -> ShortText -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ShortText -> Int -> Int -> m ()
TS.copyBytes MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
0 ShortText
str Int
0 Int
len
    MutableByteArray (PrimState (ST s))
-> Int -> Int -> Word8 -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Int -> Word8 -> m ()
fillByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
len Int
n (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
'\'')
addSuffix ShortText
str (Index Integer
n) = do
  let len :: Int
len = ShortText -> Int
TS.lengthBytes ShortText
str
      d :: Int
d = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then Int
1 else Integer -> Int
integerLog10' Integer
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  Int -> (forall s. MutableByteArray s -> ST s ()) -> ShortText
TS.unsafeCreate (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d) \MutableByteArray s
arr -> do
    MutableByteArray (PrimState (ST s))
-> Int -> ShortText -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ShortText -> Int -> Int -> m ()
TS.copyBytes MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
0 ShortText
str Int
0 Int
len
    MutableByteArray s -> Int -> Int -> Integer -> ST s ()
forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillDigits MutableByteArray s
arr (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
d Integer
n
addSuffix ShortText
str (Subscript Integer
n) = do
  let len :: Int
len = ShortText -> Int
TS.lengthBytes ShortText
str
      d :: Int
d = if Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then Int
1 else Integer -> Int
integerLog10' Integer
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  -- Every digit will require 3 bytes.
  Int -> (forall s. MutableByteArray s -> ST s ()) -> ShortText
TS.unsafeCreate (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d) \MutableByteArray s
arr -> do
    MutableByteArray (PrimState (ST s))
-> Int -> ShortText -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ShortText -> Int -> Int -> m ()
TS.copyBytes MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
0 ShortText
str Int
0 Int
len
    MutableByteArray s -> Int -> Int -> Integer -> ST s ()
forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillSubscriptDigits MutableByteArray s
arr (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
d Integer
n

-- | Fill a 'MutableByteArray' with base-10 digits.
fillDigits
  :: MutableByteArray s
  -> Int
  -- ^ Byte index of the last digit.
  -> Int
  -- ^ Number of digits.
  -> Integer
  -> ST s ()
fillDigits :: forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillDigits MutableByteArray s
arr Int
i Int
0 Integer
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
fillDigits MutableByteArray s
arr Int
i Int
ndigits Integer
n = do
  let (Integer
d, Integer
r) = Integer
n Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
10
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word8 MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
i (Integer -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
0x30 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
r))
  MutableByteArray s -> Int -> Int -> Integer -> ST s ()
forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillDigits MutableByteArray s
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
ndigits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer
d

-- | Fill a 'MutableByteArray' with base-10 subscript digits.
fillSubscriptDigits
  :: MutableByteArray s
  -> Int
  -- ^ Byte index of the last digit.
  -> Int
  -- ^ Number of digits.
  -> Integer
  -> ST s ()
fillSubscriptDigits :: forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillSubscriptDigits MutableByteArray s
arr Int
i Int
0 Integer
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
fillSubscriptDigits MutableByteArray s
arr Int
i Int
ndigits Integer
n = do
  let (Integer
d, Integer
r) = Integer
n Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
10
  -- Subscript digits ₀..₉ are encoded
  -- as 0xE2 0x82 0x80 .. 0xE2 0x82 0x89.
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word8 MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
i (Word8
0x80 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Integer -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
r)
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word8 MutableByteArray s
MutableByteArray (PrimState (ST s))
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word8
0x82
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word8 MutableByteArray s
MutableByteArray (PrimState (ST s))
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Word8
0xE2
  MutableByteArray s -> Int -> Int -> Integer -> ST s ()
forall s. MutableByteArray s -> Int -> Int -> Integer -> ST s ()
fillSubscriptDigits MutableByteArray s
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) (Int
ndigits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer
d