{-# 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
ps'nameMap :: Map Index Text
$sel:ps'nameMap:PrettyState :: PrettyState -> Map Index Text
ps'nameMap} <- forall s (m :: Type -> Type). MonadState s m => m s
get
  case Term DeBruijn uni DefaultFun ()
uplc of
    Force ()
_ (Force ()
_ (Builtin ()
_ DefaultFun
b)) -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Text
forcedPrefix forall a. Semigroup a => a -> a -> a
<> String -> Text
Txt.pack (forall a. Show a => a -> String
show DefaultFun
b)
    Force ()
_ (Builtin ()
_ DefaultFun
b) -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Text
forcedPrefix forall a. Semigroup a => a -> a -> a
<> String -> Text
Txt.pack (forall a. Show a => a -> String
show DefaultFun
b)
    Term DeBruijn uni DefaultFun ()
PFixAst -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Text
"fix"
    ComposeAST
      (Builtin () DefaultFun
PLC.SndPair)
      (Builtin () DefaultFun
PLC.UnConstrData) -> 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) -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Text
"unDataSum"
    Term DeBruijn uni DefaultFun ()
_ -> forall s. PrettyMonad s Text
freshVarName

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