{-# Language TypeOperators, DefaultSignatures, EmptyCase, GeneralizedNewtypeDeriving, MultiParamTypeClasses, FlexibleInstances, FlexibleContexts, ImportQualifiedPost, BlockArguments #-}
{-|
Module      : Advent.Tokenize
Description : Automation for replacing data with Ints
Copyright   : (c) Eric Mertens, 2022
License     : ISC
Maintainer  : emertens@gmail.com

>>> autoTokenize [("A", [("A",40,["B"])])] :: [(Int, [(Int, Integer, [Int])])]
[(0,[(0,40,[1])])]

>>> autoTokenize ("A","A","B","B") :: (Int, Int, String, Int)
(0,0,"B",1)

-}
module Advent.Tokenize where

import Control.Monad.Trans.State.Strict (State, evalState, state)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import GHC.Generics (Generic(Rep, from, to), V1, U1(U1), K1(K1), M1(M1), type (:+:)(L1,R1), type (:*:)((:*:)))

newtype Tokenize a = Tokenize (State (Map String Int) a)
  deriving ((forall a b. (a -> b) -> Tokenize a -> Tokenize b)
-> (forall a b. a -> Tokenize b -> Tokenize a) -> Functor Tokenize
forall a b. a -> Tokenize b -> Tokenize a
forall a b. (a -> b) -> Tokenize a -> Tokenize b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Tokenize a -> Tokenize b
fmap :: forall a b. (a -> b) -> Tokenize a -> Tokenize b
$c<$ :: forall a b. a -> Tokenize b -> Tokenize a
<$ :: forall a b. a -> Tokenize b -> Tokenize a
Functor, Functor Tokenize
Functor Tokenize =>
(forall a. a -> Tokenize a)
-> (forall a b. Tokenize (a -> b) -> Tokenize a -> Tokenize b)
-> (forall a b c.
    (a -> b -> c) -> Tokenize a -> Tokenize b -> Tokenize c)
-> (forall a b. Tokenize a -> Tokenize b -> Tokenize b)
-> (forall a b. Tokenize a -> Tokenize b -> Tokenize a)
-> Applicative Tokenize
forall a. a -> Tokenize a
forall a b. Tokenize a -> Tokenize b -> Tokenize a
forall a b. Tokenize a -> Tokenize b -> Tokenize b
forall a b. Tokenize (a -> b) -> Tokenize a -> Tokenize b
forall a b c.
(a -> b -> c) -> Tokenize a -> Tokenize b -> Tokenize c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> Tokenize a
pure :: forall a. a -> Tokenize a
$c<*> :: forall a b. Tokenize (a -> b) -> Tokenize a -> Tokenize b
<*> :: forall a b. Tokenize (a -> b) -> Tokenize a -> Tokenize b
$cliftA2 :: forall a b c.
(a -> b -> c) -> Tokenize a -> Tokenize b -> Tokenize c
liftA2 :: forall a b c.
(a -> b -> c) -> Tokenize a -> Tokenize b -> Tokenize c
$c*> :: forall a b. Tokenize a -> Tokenize b -> Tokenize b
*> :: forall a b. Tokenize a -> Tokenize b -> Tokenize b
$c<* :: forall a b. Tokenize a -> Tokenize b -> Tokenize a
<* :: forall a b. Tokenize a -> Tokenize b -> Tokenize a
Applicative, Applicative Tokenize
Applicative Tokenize =>
(forall a b. Tokenize a -> (a -> Tokenize b) -> Tokenize b)
-> (forall a b. Tokenize a -> Tokenize b -> Tokenize b)
-> (forall a. a -> Tokenize a)
-> Monad Tokenize
forall a. a -> Tokenize a
forall a b. Tokenize a -> Tokenize b -> Tokenize b
forall a b. Tokenize a -> (a -> Tokenize b) -> Tokenize b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. Tokenize a -> (a -> Tokenize b) -> Tokenize b
>>= :: forall a b. Tokenize a -> (a -> Tokenize b) -> Tokenize b
$c>> :: forall a b. Tokenize a -> Tokenize b -> Tokenize b
>> :: forall a b. Tokenize a -> Tokenize b -> Tokenize b
$creturn :: forall a. a -> Tokenize a
return :: forall a. a -> Tokenize a
Monad)

tok :: String -> Tokenize Int
tok :: String -> Tokenize Int
tok String
t = State (Map String Int) Int -> Tokenize Int
forall a. State (Map String Int) a -> Tokenize a
Tokenize (State (Map String Int) Int -> Tokenize Int)
-> State (Map String Int) Int -> Tokenize Int
forall a b. (a -> b) -> a -> b
$
    (Map String Int -> (Int, Map String Int))
-> State (Map String Int) Int
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state \Map String Int
m ->
        case String -> Map String Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
t Map String Int
m of
            Just Int
i -> (Int
i,Map String Int
m)
            Maybe Int
Nothing -> Int
i Int -> (Int, Map String Int) -> (Int, Map String Int)
forall a b. a -> b -> b
`seq` Map String Int
m' Map String Int -> (Int, Map String Int) -> (Int, Map String Int)
forall a b. a -> b -> b
`seq` (Int
i,Map String Int
m')
                where
                    i :: Int
i = Map String Int -> Int
forall k a. Map k a -> Int
Map.size Map String Int
m
                    m' :: Map String Int
m' = String -> Int -> Map String Int -> Map String Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
t Int
i Map String Int
m

runTokenize :: Tokenize a -> a
runTokenize :: forall a. Tokenize a -> a
runTokenize (Tokenize State (Map String Int) a
x) = State (Map String Int) a -> Map String Int -> a
forall s a. State s a -> s -> a
evalState State (Map String Int) a
x Map String Int
forall k a. Map k a
Map.empty

autoTokenize :: AutoToken a b => a -> b
autoTokenize :: forall a b. AutoToken a b => a -> b
autoTokenize = Tokenize b -> b
forall a. Tokenize a -> a
runTokenize (Tokenize b -> b) -> (a -> Tokenize b) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Tokenize b
forall a b. AutoToken a b => a -> Tokenize b
autoToken

class AutoToken a b where
    autoToken :: a -> Tokenize b
    default autoToken :: (Generic a, Generic b, GAutoToken (Rep a) (Rep b)) => a -> Tokenize b
    autoToken = a -> Tokenize b
forall a b.
(Generic a, Generic b, GAutoToken (Rep a) (Rep b)) =>
a -> Tokenize b
genericAutoToken

instance {-# INCOHERENT #-} AutoToken a a where
    autoToken :: a -> Tokenize a
autoToken = a -> Tokenize a
forall a. a -> Tokenize a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance AutoToken String Int where
    autoToken :: String -> Tokenize Int
autoToken = String -> Tokenize Int
tok

instance AutoToken a b => AutoToken (Map String a) (IntMap b) where
    autoToken :: Map String a -> Tokenize (IntMap b)
autoToken Map String a
m = [(Int, b)] -> IntMap b
forall a. [(Int, a)] -> IntMap a
IntMap.fromList ([(Int, b)] -> IntMap b)
-> Tokenize [(Int, b)] -> Tokenize (IntMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(String, a)] -> Tokenize [(Int, b)]
forall a b. AutoToken a b => a -> Tokenize b
autoToken (Map String a -> [(String, a)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map String a
m)

instance (Ord b, AutoToken a b) => AutoToken (Set a) (Set b) where
    autoToken :: Set a -> Tokenize (Set b)
autoToken Set a
m = [b] -> Set b
forall a. Ord a => [a] -> Set a
Set.fromList ([b] -> Set b) -> Tokenize [b] -> Tokenize (Set b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Tokenize [b]
forall a b. AutoToken a b => a -> Tokenize b
autoToken (Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
m)

instance AutoToken (Set String) IntSet where
    autoToken :: Set String -> Tokenize IntSet
autoToken Set String
m = [Int] -> IntSet
IntSet.fromList ([Int] -> IntSet) -> Tokenize [Int] -> Tokenize IntSet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String] -> Tokenize [Int]
forall a b. AutoToken a b => a -> Tokenize b
autoToken (Set String -> [String]
forall a. Set a -> [a]
Set.toList Set String
m)

instance AutoToken a b => AutoToken [a] [b] where
    autoToken :: [a] -> Tokenize [b]
autoToken = (a -> Tokenize b) -> [a] -> Tokenize [b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse a -> Tokenize b
forall a b. AutoToken a b => a -> Tokenize b
autoToken

instance AutoToken a b => AutoToken (Maybe a) (Maybe b) where
    autoToken :: Maybe a -> Tokenize (Maybe b)
autoToken = (a -> Tokenize b) -> Maybe a -> Tokenize (Maybe b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse a -> Tokenize b
forall a b. AutoToken a b => a -> Tokenize b
autoToken

instance (AutoToken a1 b1, AutoToken a2 b2) => AutoToken (a1,a2) (b1,b2)

instance (AutoToken a1 b1, AutoToken a2 b2, AutoToken a3 b3) =>
    AutoToken (a1,a2,a3) (b1,b2,b3)

instance (AutoToken a1 b1, AutoToken a2 b2, AutoToken a3 b3, AutoToken a4 b4) =>
    AutoToken (a1,a2,a3,a4) (b1,b2,b3,b4)


class GAutoToken f g where
    gautoToken :: f x -> Tokenize (g x)

instance GAutoToken f g => GAutoToken (M1 i c f) (M1 j d g) where
    gautoToken :: forall x. M1 i c f x -> Tokenize (M1 j d g x)
gautoToken (M1 f x
x) = g x -> M1 j d g x
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 (g x -> M1 j d g x) -> Tokenize (g x) -> Tokenize (M1 j d g x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f x -> Tokenize (g x)
forall x. f x -> Tokenize (g x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken f x
x

instance (GAutoToken f1 f2, GAutoToken g1 g2) =>
    GAutoToken (f1 :*: g1) (f2 :*: g2) where
    gautoToken :: forall x. (:*:) f1 g1 x -> Tokenize ((:*:) f2 g2 x)
gautoToken (f1 x
x :*: g1 x
y) = f2 x -> g2 x -> (:*:) f2 g2 x
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) (f2 x -> g2 x -> (:*:) f2 g2 x)
-> Tokenize (f2 x) -> Tokenize (g2 x -> (:*:) f2 g2 x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f1 x -> Tokenize (f2 x)
forall x. f1 x -> Tokenize (f2 x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken f1 x
x Tokenize (g2 x -> (:*:) f2 g2 x)
-> Tokenize (g2 x) -> Tokenize ((:*:) f2 g2 x)
forall a b. Tokenize (a -> b) -> Tokenize a -> Tokenize b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> g1 x -> Tokenize (g2 x)
forall x. g1 x -> Tokenize (g2 x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken g1 x
y

instance (GAutoToken f1 f2, GAutoToken g1 g2) =>
    GAutoToken (f1 :+: g1) (f2 :+: g2) where
    gautoToken :: forall x. (:+:) f1 g1 x -> Tokenize ((:+:) f2 g2 x)
gautoToken (L1 f1 x
x) = f2 x -> (:+:) f2 g2 x
forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 (f2 x -> (:+:) f2 g2 x)
-> Tokenize (f2 x) -> Tokenize ((:+:) f2 g2 x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f1 x -> Tokenize (f2 x)
forall x. f1 x -> Tokenize (f2 x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken f1 x
x
    gautoToken (R1 g1 x
x) = g2 x -> (:+:) f2 g2 x
forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 (g2 x -> (:+:) f2 g2 x)
-> Tokenize (g2 x) -> Tokenize ((:+:) f2 g2 x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g1 x -> Tokenize (g2 x)
forall x. g1 x -> Tokenize (g2 x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken g1 x
x

instance AutoToken a b => GAutoToken (K1 i a) (K1 i b) where
    gautoToken :: forall x. K1 i a x -> Tokenize (K1 i b x)
gautoToken (K1 a
x) = b -> K1 i b x
forall k i c (p :: k). c -> K1 i c p
K1 (b -> K1 i b x) -> Tokenize b -> Tokenize (K1 i b x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Tokenize b
forall a b. AutoToken a b => a -> Tokenize b
autoToken a
x

instance GAutoToken U1 U1 where
    gautoToken :: forall x. U1 x -> Tokenize (U1 x)
gautoToken U1 x
_ = U1 x -> Tokenize (U1 x)
forall a. a -> Tokenize a
forall (f :: * -> *) a. Applicative f => a -> f a
pure U1 x
forall k (p :: k). U1 p
U1

instance GAutoToken V1 V1 where
    gautoToken :: forall x. V1 x -> Tokenize (V1 x)
gautoToken V1 x
v = case V1 x
v of {}

genericAutoToken ::
    (Generic a, Generic b, GAutoToken (Rep a) (Rep b)) =>
    a -> Tokenize b
genericAutoToken :: forall a b.
(Generic a, Generic b, GAutoToken (Rep a) (Rep b)) =>
a -> Tokenize b
genericAutoToken a
x = Rep b Any -> b
forall a x. Generic a => Rep a x -> a
forall x. Rep b x -> b
to (Rep b Any -> b) -> Tokenize (Rep b Any) -> Tokenize b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rep a Any -> Tokenize (Rep b Any)
forall x. Rep a x -> Tokenize (Rep b x)
forall (f :: * -> *) (g :: * -> *) x.
GAutoToken f g =>
f x -> Tokenize (g x)
gautoToken (a -> Rep a Any
forall x. a -> Rep a x
forall a x. Generic a => a -> Rep a x
from a
x)
{-# INLINE genericAutoToken #-}