{- |
Module      : Control.Lens.Grammar.Token
Description : lexical tokens
Copyright   : (C) 2026 - Eitan Chatav
License     : BSD-style (see the file LICENSE)
Maintainer  : Eitan Chatav <eitan.chatav@gmail.com>
Stability   : provisional
Portability : non-portable
-}

module Control.Lens.Grammar.Token
  ( -- * Tokenized
    Tokenized (..)
  , satisfy
  , tokens
    -- * Categorized
  , Categorized (..)
  , GeneralCategory (..)
  ) where

import Control.Lens
import Control.Lens.PartialIso
import Control.Monad.Loops (iterateUntil)
import Data.Bifunctor.Joker
import Data.Char
import Data.Foldable
import Data.Profunctor
import Data.Profunctor.Monoidal
import Data.Word
import Control.Monad.State (StateT, state)
import System.Random (RandomGen, Random, random, randomR)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Gen (Gen)
import qualified Test.QuickCheck.Gen as Gen
import Text.ParserCombinators.ReadP (ReadP)
import qualified Text.ParserCombinators.ReadP as ReadP

{- | `Categorized` provides a type family `Categorize`
and a function to `categorize` tokens into disjoint categories.

>>> :kind! Categorize Char
Categorize Char :: *
= GeneralCategory

>>> categorize 'a'
LowercaseLetter
-}
class (Ord token, Ord (Categorize token), Enum (Categorize token))
  => Categorized token where
  type Categorize token
  type Categorize token = ()
  categorize :: token -> Categorize token
  default categorize :: Categorize token ~ () => token -> Categorize token
  categorize token
_ = ()
instance Categorized Char where
  type Categorize Char = GeneralCategory
  categorize :: Char -> Categorize Char
categorize = Char -> GeneralCategory
Char -> Categorize Char
generalCategory
instance Categorized Word8
instance Categorized ()

{- | `Tokenized` combinators for constructing lexical tokens. -}
class Categorized token => Tokenized token p | p -> token where
  {- | Any single token. -}
  anyToken :: p

  {- | A single specified `token`. -}
  token :: token -> p
  default token
    :: (p ~ q token token, Choice q, Cochoice q)
    => token -> p
  token = (token -> Bool) -> q token token
forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy ((token -> Bool) -> q token token)
-> (token -> token -> Bool) -> token -> q token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. token -> token -> Bool
forall token p. Tokenized token p => token -> p
token

  {- | A single token which is `oneOf` a set.

  prop> token x = oneOf [x]
  
  -}
  oneOf :: Foldable f => f token -> p
  default oneOf
    :: (p ~ q token token, Choice q, Cochoice q, Foldable f)
    => f token -> p
  oneOf = (token -> Bool) -> q token token
forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy ((token -> Bool) -> q token token)
-> (f token -> token -> Bool) -> f token -> q token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f token -> token -> Bool
forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
forall (f :: * -> *). Foldable f => f token -> token -> Bool
oneOf

  {- | A single token which is `notOneOf` a set.
  
  prop> anyToken = notOneOf []
  
  -}
  notOneOf :: Foldable f => f token -> p
  default notOneOf
    :: (p ~ q token token, Choice q, Cochoice q, Foldable f)
    => f token -> p
  notOneOf = (token -> Bool) -> q token token
forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy ((token -> Bool) -> q token token)
-> (f token -> token -> Bool) -> f token -> q token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f token -> token -> Bool
forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
forall (f :: * -> *). Foldable f => f token -> token -> Bool
notOneOf

  {- | A single token which is `asIn` a category. -}
  asIn :: Categorize token -> p
  default asIn
    :: (p ~ q token token, Choice q, Cochoice q)
    => Categorize token -> p
  asIn = (token -> Bool) -> q token token
forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy ((token -> Bool) -> q token token)
-> (Categorize token -> token -> Bool)
-> Categorize token
-> q token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize token -> token -> Bool
forall token p. Tokenized token p => Categorize token -> p
asIn

  {- | A single token which is `notAsIn` a category. -}
  notAsIn :: Categorize token -> p
  default notAsIn
    :: (p ~ q token token, Choice q, Cochoice q)
    => Categorize token -> p
  notAsIn = (token -> Bool) -> q token token
forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy ((token -> Bool) -> q token token)
-> (Categorize token -> token -> Bool)
-> Categorize token
-> q token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize token -> token -> Bool
forall token p. Tokenized token p => Categorize token -> p
notAsIn

{- | A single token that satisfies a predicate. -}
satisfy
  :: (Tokenized a (p a a), Choice p, Cochoice p)
  => (a -> Bool) -> p a a
satisfy :: forall a (p :: * -> * -> *).
(Tokenized a (p a a), Choice p, Cochoice p) =>
(a -> Bool) -> p a a
satisfy a -> Bool
f = (a -> Bool) -> PartialIso' a a
forall a. (a -> Bool) -> PartialIso' a a
satisfied a -> Bool
f (PartialExchange a a a (Maybe a)
 -> PartialExchange a a a (Maybe a))
-> p a a -> p a a
forall (p :: * -> * -> *) s t a b.
(Choice p, Cochoice p) =>
APartialIso s t a b -> p a b -> p s t
>?< p a a
forall token p. Tokenized token p => p
anyToken

{- | A specified stream of `tokens`. -}
tokens
  :: ( Foldable f, Tokenized a (p a a)
     , Monoidal p, Choice p
     , AsEmpty s, Cons s s a a
     )
  => f a -> p s s
tokens :: forall (f :: * -> *) a (p :: * -> * -> *) s.
(Foldable f, Tokenized a (p a a), Monoidal p, Choice p, AsEmpty s,
 Cons s s a a) =>
f a -> p s s
tokens = (a -> p s s -> p s s) -> p s s -> f a -> p s s
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (p a a -> p s s -> p s s
forall s t a b (p :: * -> * -> *).
(Cons s t a b, Monoidal p, Choice p) =>
p a b -> p s t -> p s t
(>:<) (p a a -> p s s -> p s s) -> (a -> p a a) -> a -> p s s -> p s s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> p a a
forall token p. Tokenized token p => token -> p
token) p s s
forall s (p :: * -> * -> *).
(AsEmpty s, Monoidal p, Choice p) =>
p s s
asEmpty

-- instances
instance Categorized token => Tokenized token (token -> Bool) where
  anyToken :: token -> Bool
anyToken token
_ = Bool
True
  token :: token -> token -> Bool
token = token -> token -> Bool
forall a. Eq a => a -> a -> Bool
(==)
  oneOf :: forall (f :: * -> *). Foldable f => f token -> token -> Bool
oneOf = (token -> f token -> Bool) -> f token -> token -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip token -> f token -> Bool
forall a. Eq a => a -> f a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem
  notOneOf :: forall (f :: * -> *). Foldable f => f token -> token -> Bool
notOneOf = (token -> f token -> Bool) -> f token -> token -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip token -> f token -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
notElem
  asIn :: Categorize token -> token -> Bool
asIn = (token -> Categorize token)
-> (Categorize token -> Bool) -> token -> Bool
forall a b c. (a -> b) -> (b -> c) -> a -> c
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize ((Categorize token -> Bool) -> token -> Bool)
-> (Categorize token -> Categorize token -> Bool)
-> Categorize token
-> token
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
(==)
  notAsIn :: Categorize token -> token -> Bool
notAsIn = (token -> Categorize token)
-> (Categorize token -> Bool) -> token -> Bool
forall a b c. (a -> b) -> (b -> c) -> a -> c
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize ((Categorize token -> Bool) -> token -> Bool)
-> (Categorize token -> Categorize token -> Bool)
-> Categorize token
-> token
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
(/=)
instance Tokenized token (f token)
  => Tokenized token (Joker f token token) where
    anyToken :: Joker f token token
anyToken = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (forall token p. Tokenized token p => p
anyToken @token)
    token :: token -> Joker f token token
token = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (f token -> Joker f token token)
-> (token -> f token) -> token -> Joker f token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall token p. Tokenized token p => token -> p
token @token
    oneOf :: forall (f :: * -> *). Foldable f => f token -> Joker f token token
oneOf = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (f token -> Joker f token token)
-> (f token -> f token) -> f token -> Joker f token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
oneOf @token
    notOneOf :: forall (f :: * -> *). Foldable f => f token -> Joker f token token
notOneOf = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (f token -> Joker f token token)
-> (f token -> f token) -> f token -> Joker f token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
notOneOf @token
    asIn :: Categorize token -> Joker f token token
asIn = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (f token -> Joker f token token)
-> (Categorize token -> f token)
-> Categorize token
-> Joker f token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall token p. Tokenized token p => Categorize token -> p
asIn @token
    notAsIn :: Categorize token -> Joker f token token
notAsIn = f token -> Joker f token token
forall {k} {k1} (g :: k -> *) (a :: k1) (b :: k).
g b -> Joker g a b
Joker (f token -> Joker f token token)
-> (Categorize token -> f token)
-> Categorize token
-> Joker f token token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall token p. Tokenized token p => Categorize token -> p
notAsIn @token
instance Tokenized Char (ReadP Char) where
  anyToken :: ReadP Char
anyToken = ReadP Char
ReadP.get
  token :: Char -> ReadP Char
token = Char -> ReadP Char
ReadP.char
  oneOf :: forall (f :: * -> *). Foldable f => f Char -> ReadP Char
oneOf = (Char -> Bool) -> ReadP Char
ReadP.satisfy ((Char -> Bool) -> ReadP Char)
-> (f Char -> Char -> Bool) -> f Char -> ReadP Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Char -> Char -> Bool
forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
forall (f :: * -> *). Foldable f => f Char -> Char -> Bool
oneOf
  notOneOf :: forall (f :: * -> *). Foldable f => f Char -> ReadP Char
notOneOf = (Char -> Bool) -> ReadP Char
ReadP.satisfy ((Char -> Bool) -> ReadP Char)
-> (f Char -> Char -> Bool) -> f Char -> ReadP Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Char -> Char -> Bool
forall token p (f :: * -> *).
(Tokenized token p, Foldable f) =>
f token -> p
forall (f :: * -> *). Foldable f => f Char -> Char -> Bool
notOneOf
  asIn :: Categorize Char -> ReadP Char
asIn = (Char -> Bool) -> ReadP Char
ReadP.satisfy ((Char -> Bool) -> ReadP Char)
-> (Categorize Char -> Char -> Bool)
-> Categorize Char
-> ReadP Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize Char -> Char -> Bool
forall token p. Tokenized token p => Categorize token -> p
asIn
  notAsIn :: Categorize Char -> ReadP Char
notAsIn = (Char -> Bool) -> ReadP Char
ReadP.satisfy ((Char -> Bool) -> ReadP Char)
-> (Categorize Char -> Char -> Bool)
-> Categorize Char
-> ReadP Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Categorize Char -> Char -> Bool
forall token p. Tokenized token p => Categorize token -> p
notAsIn
instance (Categorized token, Arbitrary token) => Tokenized token (Gen token) where
  anyToken :: Gen token
anyToken = forall a. Arbitrary a => Gen a
arbitrary @token
  token :: token -> Gen token
token = token -> Gen token
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  oneOf :: forall (f :: * -> *). Foldable f => f token -> Gen token
oneOf = [token] -> Gen token
forall a. HasCallStack => [a] -> Gen a
Gen.elements ([token] -> Gen token)
-> (f token -> [token]) -> f token -> Gen token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f token -> [token]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
  notOneOf :: forall (f :: * -> *). Foldable f => f token -> Gen token
notOneOf f token
xs = Gen token
forall a. Arbitrary a => Gen a
arbitrary Gen token -> (token -> Bool) -> Gen token
forall a. Gen a -> (a -> Bool) -> Gen a
`Gen.suchThat` (token -> f token -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` f token
xs)
  asIn :: Categorize token -> Gen token
asIn Categorize token
cat = Gen token
forall a. Arbitrary a => Gen a
arbitrary Gen token -> (token -> Bool) -> Gen token
forall a. Gen a -> (a -> Bool) -> Gen a
`Gen.suchThat` (Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
(==) Categorize token
cat (Categorize token -> Bool)
-> (token -> Categorize token) -> token -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize)
  notAsIn :: Categorize token -> Gen token
notAsIn Categorize token
cat = Gen token
forall a. Arbitrary a => Gen a
arbitrary Gen token -> (token -> Bool) -> Gen token
forall a. Gen a -> (a -> Bool) -> Gen a
`Gen.suchThat` (Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
(/=) Categorize token
cat (Categorize token -> Bool)
-> (token -> Categorize token) -> token -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize)
instance (RandomGen g, Monad m, Categorized token, Random token)
  => Tokenized token (StateT g m token) where
  anyToken :: StateT g m token
anyToken = (g -> (token, g)) -> StateT g m token
forall a. (g -> (a, g)) -> StateT g m a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state g -> (token, g)
forall g. RandomGen g => g -> (token, g)
forall a g. (Random a, RandomGen g) => g -> (a, g)
random
  token :: token -> StateT g m token
token = token -> StateT g m token
forall a. a -> StateT g m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  oneOf :: forall (f :: * -> *). Foldable f => f token -> StateT g m token
oneOf f token
xs = do
    let ys :: [token]
ys = f token -> [token]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f token
xs
    Int
i <- (g -> (Int, g)) -> StateT g m Int
forall a. (g -> (a, g)) -> StateT g m a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((Int, Int) -> g -> (Int, g)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
0, [token] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [token]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
    token -> StateT g m token
forall a. a -> StateT g m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([token]
ys [token] -> Int -> token
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
  notOneOf :: forall (f :: * -> *). Foldable f => f token -> StateT g m token
notOneOf f token
xs = (token -> Bool) -> StateT g m token -> StateT g m token
forall (m :: * -> *) a. Monad m => (a -> Bool) -> m a -> m a
iterateUntil (token -> f token -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` f token
xs) StateT g m token
forall token p. Tokenized token p => p
anyToken
  asIn :: Categorize token -> StateT g m token
asIn Categorize token
cat = (token -> Bool) -> StateT g m token -> StateT g m token
forall (m :: * -> *) a. Monad m => (a -> Bool) -> m a -> m a
iterateUntil ((Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
== Categorize token
cat) (Categorize token -> Bool)
-> (token -> Categorize token) -> token -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize) StateT g m token
forall token p. Tokenized token p => p
anyToken
  notAsIn :: Categorize token -> StateT g m token
notAsIn Categorize token
cat = (token -> Bool) -> StateT g m token -> StateT g m token
forall (m :: * -> *) a. Monad m => (a -> Bool) -> m a -> m a
iterateUntil ((Categorize token -> Categorize token -> Bool
forall a. Eq a => a -> a -> Bool
/= Categorize token
cat) (Categorize token -> Bool)
-> (token -> Categorize token) -> token -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. token -> Categorize token
forall token. Categorized token => token -> Categorize token
categorize) StateT g m token
forall token p. Tokenized token p => p
anyToken