{-# LANGUAGE PatternSynonyms #-}

module Plutarch.Pretty.Internal.Name (smartName, freshVarName) where

import Control.Monad.Reader (ask)
import Control.Monad.State (
  get,
  lift,
  modify',
 )
import Data.Functor (($>))
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Txt
import Data.Traversable (for)

import System.Random.Stateful (randomRM, uniformRM)

import PlutusCore qualified as PLC
import UntypedPlutusCore (
  DeBruijn (DeBruijn),
  DefaultFun,
  Term (Builtin, Force, Var),
 )

import Plutarch.Pretty.Internal.Config (forcedPrefix, keywords)
import Plutarch.Pretty.Internal.TermUtils (pattern ComposeAST, pattern PFixAst)
import Plutarch.Pretty.Internal.Types (
  PrettyMonad,
  PrettyState (PrettyState, ps'nameMap, ps'names),
  builtinFunAtRef,
  memorizeName,
 )

smartName :: Term DeBruijn uni DefaultFun () -> PrettyMonad s Text
smartName :: forall (uni :: Type -> Type) s.
Term DeBruijn uni DefaultFun () -> PrettyMonad s Text
smartName Term DeBruijn uni DefaultFun ()
uplc = do
  PrettyState {Map Index Text
$sel:ps'nameMap:PrettyState :: PrettyState -> Map Index Text
ps'nameMap :: Map Index Text
ps'nameMap} <- ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) PrettyState
forall s (m :: Type -> Type). MonadState s m => m s
get
  case Term DeBruijn uni DefaultFun ()
uplc of
    Force ()
_ (Force ()
_ (Builtin ()
_ DefaultFun
b)) -> Text -> PrettyMonad s Text
forall a.
a -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Text -> PrettyMonad s Text) -> Text -> PrettyMonad s Text
forall a b. (a -> b) -> a -> b
$ Text
forcedPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Txt.pack (DefaultFun -> String
forall a. Show a => a -> String
show DefaultFun
b)
    Force ()
_ (Builtin ()
_ DefaultFun
b) -> Text -> PrettyMonad s Text
forall a.
a -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Text -> PrettyMonad s Text) -> Text -> PrettyMonad s Text
forall a b. (a -> b) -> a -> b
$ Text
forcedPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Txt.pack (DefaultFun -> String
forall a. Show a => a -> String
show DefaultFun
b)
    Term DeBruijn uni DefaultFun ()
PFixAst -> Text -> PrettyMonad s Text
forall a.
a -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Text
"fix"
    ComposeAST
      (Builtin () DefaultFun
PLC.SndPair)
      (Builtin () DefaultFun
PLC.UnConstrData) -> Text -> PrettyMonad s Text
forall a.
a -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Text
"unDataSum"
    ComposeAST
      (Var () (DeBruijn (Map Index Text -> Index -> Maybe DefaultFun
builtinFunAtRef Map Index Text
ps'nameMap -> Just DefaultFun
PLC.SndPair)))
      (Builtin () DefaultFun
PLC.UnConstrData) -> Text -> PrettyMonad s Text
forall a.
a -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Text
"unDataSum"
    Term DeBruijn uni DefaultFun ()
_ -> PrettyMonad s Text
forall s. PrettyMonad s Text
freshVarName

freshVarName :: PrettyMonad s Text
freshVarName :: forall s. PrettyMonad s Text
freshVarName = do
  STGenM StdGen s
stGen <- ReaderT
  (STGenM StdGen s) (StateT PrettyState (ST s)) (STGenM StdGen s)
forall r (m :: Type -> Type). MonadReader r m => m r
ask
  PrettyState {Set Text
$sel:ps'names:PrettyState :: PrettyState -> Set Text
ps'names :: Set Text
ps'names} <- ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) PrettyState
forall s (m :: Type -> Type). MonadState s m => m s
get
  let existingNames :: Set Text
existingNames = Set Text -> Set Text -> Set Text
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set Text
ps'names Set Text
keywords
  Int
nameTailLen <- StateT PrettyState (ST s) Int
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Int
forall (m :: Type -> Type) a.
Monad m =>
m a -> ReaderT (STGenM StdGen s) m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT PrettyState (ST s) Int
 -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Int)
-> (ST s Int -> StateT PrettyState (ST s) Int)
-> ST s Int
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST s Int -> StateT PrettyState (ST s) Int
forall (m :: Type -> Type) a.
Monad m =>
m a -> StateT PrettyState m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ST s Int
 -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Int)
-> ST s Int
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Int
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> STGenM StdGen s -> ST s Int
forall g r (m :: Type -> Type) a.
(RandomGenM g r m, Random a) =>
(a, a) -> g -> m a
randomRM (Int
0 :: Int, Int
7) STGenM StdGen s
stGen
  Char
beginChar <- Text -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
forall {a} {t :: (Type -> Type) -> Type -> Type}
       {t :: (Type -> Type) -> Type -> Type} {m :: Type -> Type}.
(MonadReader a (t (t m)), MonadTrans t, MonadTrans t, Monad (t m),
 StatefulGen a m) =>
Text -> t (t m) Char
chooseChar Text
starterChars
  Text
newName <- (String -> Text)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String
-> PrettyMonad s Text
forall a b.
(a -> b)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) a
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> Text
Txt.pack (String -> Text) -> (String -> String) -> String -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
beginChar :)) (ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String
 -> PrettyMonad s Text)
-> (ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
    -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
-> PrettyMonad s Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int]
-> (Int
    -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Int
Item [Int]
0 .. Int
Item [Int]
nameTailLen] ((Int
  -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char)
 -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String)
-> (ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
    -> Int
    -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
-> Int
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
forall a b. a -> b -> a
const (ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
 -> PrettyMonad s Text)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
-> PrettyMonad s Text
forall a b. (a -> b) -> a -> b
$ Text -> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) Char
forall {a} {t :: (Type -> Type) -> Type -> Type}
       {t :: (Type -> Type) -> Type -> Type} {m :: Type -> Type}.
(MonadReader a (t (t m)), MonadTrans t, MonadTrans t, Monad (t m),
 StatefulGen a m) =>
Text -> t (t m) Char
chooseChar Text
chars
  if Text -> Set Text -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Text
newName Set Text
existingNames
    then PrettyMonad s Text
forall s. PrettyMonad s Text
freshVarName
    else (PrettyState -> PrettyState)
-> ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify' (Text -> PrettyState -> PrettyState
memorizeName Text
newName) ReaderT (STGenM StdGen s) (StateT PrettyState (ST s)) ()
-> Text -> PrettyMonad s Text
forall (f :: Type -> Type) a b. Functor f => f a -> b -> f b
$> Text
newName
  where
    chooseChar :: Text -> t (t m) Char
chooseChar Text
x = do
      a
stGen <- t (t m) a
forall r (m :: Type -> Type). MonadReader r m => m r
ask
      Int
chosenIx <- t m Int -> t (t m) Int
forall (m :: Type -> Type) a. Monad m => m a -> t m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (t m Int -> t (t m) Int)
-> (m Int -> t m Int) -> m Int -> t (t m) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m Int -> t m Int
forall (m :: Type -> Type) a. Monad m => m a -> t m a
forall (t :: (Type -> Type) -> Type -> Type) (m :: Type -> Type) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Int -> t (t m) Int) -> m Int -> t (t m) Int
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> a -> m Int
forall a g (m :: Type -> Type).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: Type -> Type).
StatefulGen g m =>
(Int, Int) -> g -> m Int
uniformRM (Int
0, Text -> Int
Txt.length Text
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
stGen
      Char -> t (t m) Char
forall a. a -> t (t m) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Char -> t (t m) Char) -> Char -> t (t m) Char
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Int -> Char
Text -> Int -> Char
Txt.index Text
x Int
chosenIx
    starterChars :: Text
starterChars = String -> Text
Txt.pack [Char
Item String
'a' .. Char
Item String
'z']
    chars :: Text
chars = Text -> Text -> Text
Txt.append Text
starterChars (Text -> Text) -> (String -> Text) -> String -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Txt.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ ([Char
Item String
'A' .. Char
Item String
'Z'] String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ([Char
Item String
'0' .. Char
Item String
'9'] String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [Char
Item String
'_', Char
Item String
'\'']))