{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}

-- | Functions for compiling GHC names into Plutus Core names.
module PlutusTx.Compiler.Names where

import PlutusTx.Compiler.Kind
import {-# SOURCE #-} PlutusTx.Compiler.Type
import PlutusTx.Compiler.Types
import PlutusTx.PIRTypes
import PlutusTx.PLCTypes

import GHC.Plugins qualified as GHC

import PlutusCore qualified as PLC
import PlutusCore.MkPlc qualified as PLC
import PlutusCore.Quote

import PlutusIR.Compiler.Names

import Data.Char
import Data.Functor
import Data.List
import Data.Map qualified as Map
import Data.Text qualified as T

lookupName :: Scope uni -> GHC.Name -> Maybe (PLCVar uni)
lookupName :: forall (uni :: * -> *). Scope uni -> Name -> Maybe (PLCVar uni)
lookupName (Scope Map Name (PLCVar uni)
ns Map Name PLCTyVar
_) Name
n = Name -> Map Name (PLCVar uni) -> Maybe (PLCVar uni)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name
n Map Name (PLCVar uni)
ns

{- |
Reverses the OccName tidying that GHC does, see 'tidyOccEnv'
and accompanying Notes.

This is bad, because it makes it much harder to read since the
disambiguating numbers are gone. However, these appear to be
non-deterministic (possibly depending on the order in which
modules are processed?), so we can't rely on them.

Essentially, we just strip off trailing digits.
This might remove "real" digits added by the user, but
there's not much we can do about that.

Note that this only affects the *textual* name, not the underlying
unique, so it has no effect on the behaviour of the program, merely
on how it is printed.
-}
getUntidiedOccString :: GHC.Name -> String
getUntidiedOccString :: Name -> String
getUntidiedOccString Name
n = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Char -> Bool
isDigit (Name -> String
forall a. NamedThing a => a -> String
GHC.getOccString Name
n)

compileNameFresh :: (MonadQuote m) => GHC.Name -> m PLC.Name
compileNameFresh :: forall (m :: * -> *). MonadQuote m => Name -> m Name
compileNameFresh Name
n = Text -> m Name
forall (m :: * -> *). MonadQuote m => Text -> m Name
safeFreshName (Text -> m Name) -> Text -> m Name
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Name -> String
getUntidiedOccString Name
n

compileVarFresh :: (CompilingDefault uni fun m ann) => Ann -> GHC.Var -> m (PLCVar uni)
compileVarFresh :: forall (uni :: * -> *) fun (m :: * -> *) ann.
CompilingDefault uni fun m ann =>
Ann -> Var -> m (PLCVar uni)
compileVarFresh Ann
ann Var
v = do
  PIRType uni
t' <- Type -> m (PIRType uni)
forall (uni :: * -> *) fun (m :: * -> *) ann.
CompilingDefault uni fun m ann =>
Type -> m (PIRType uni)
compileTypeNorm (Type -> m (PIRType uni)) -> Type -> m (PIRType uni)
forall a b. (a -> b) -> a -> b
$ Var -> Type
GHC.varType Var
v
  Name
n' <- Name -> m Name
forall (m :: * -> *). MonadQuote m => Name -> m Name
compileNameFresh (Name -> m Name) -> Name -> m Name
forall a b. (a -> b) -> a -> b
$ Var -> Name
forall a. NamedThing a => a -> Name
GHC.getName Var
v
  PLCVar uni -> m (PLCVar uni)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PLCVar uni -> m (PLCVar uni)) -> PLCVar uni -> m (PLCVar uni)
forall a b. (a -> b) -> a -> b
$ Ann -> Name -> PIRType uni -> PLCVar uni
forall tyname name (uni :: * -> *) ann.
ann -> name -> Type tyname uni ann -> VarDecl tyname name uni ann
PLC.VarDecl Ann
ann Name
n' PIRType uni
t'

{- | Like `compileVarFresh`, but takes a `PIRType` instead of obtaining the
PIR type from the given `GHC.Var`.
-}
compileVarWithTyFresh ::
  (CompilingDefault uni fun m ann) =>
  Ann ->
  GHC.Var ->
  PIRType uni ->
  m (PLCVar uni)
compileVarWithTyFresh :: forall (uni :: * -> *) fun (m :: * -> *) ann.
CompilingDefault uni fun m ann =>
Ann -> Var -> PIRType uni -> m (PLCVar uni)
compileVarWithTyFresh Ann
ann Var
v PIRType uni
t = do
  Name
n' <- Name -> m Name
forall (m :: * -> *). MonadQuote m => Name -> m Name
compileNameFresh (Name -> m Name) -> Name -> m Name
forall a b. (a -> b) -> a -> b
$ Var -> Name
forall a. NamedThing a => a -> Name
GHC.getName Var
v
  PLCVar uni -> m (PLCVar uni)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PLCVar uni -> m (PLCVar uni)) -> PLCVar uni -> m (PLCVar uni)
forall a b. (a -> b) -> a -> b
$ Ann -> Name -> PIRType uni -> PLCVar uni
forall tyname name (uni :: * -> *) ann.
ann -> name -> Type tyname uni ann -> VarDecl tyname name uni ann
PLC.VarDecl Ann
ann Name
n' PIRType uni
t

lookupTyName :: Scope uni -> GHC.Name -> Maybe PLCTyVar
lookupTyName :: forall (uni :: * -> *). Scope uni -> Name -> Maybe PLCTyVar
lookupTyName (Scope Map Name (PLCVar uni)
_ Map Name PLCTyVar
tyns) Name
n = Name -> Map Name PLCTyVar -> Maybe PLCTyVar
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name
n Map Name PLCTyVar
tyns

compileTyNameFresh :: (MonadQuote m) => GHC.Name -> m PLC.TyName
compileTyNameFresh :: forall (m :: * -> *). MonadQuote m => Name -> m TyName
compileTyNameFresh Name
n = Text -> m TyName
forall (m :: * -> *). MonadQuote m => Text -> m TyName
safeFreshTyName (Text -> m TyName) -> Text -> m TyName
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Name -> String
getUntidiedOccString Name
n

compileTyVarFresh :: (Compiling uni fun m ann) => GHC.TyVar -> m PLCTyVar
compileTyVarFresh :: forall (uni :: * -> *) fun (m :: * -> *) ann.
Compiling uni fun m ann =>
Var -> m PLCTyVar
compileTyVarFresh Var
v = do
  Kind ()
k' <- Type -> m (Kind ())
forall (uni :: * -> *) fun (m :: * -> *) ann.
Compiling uni fun m ann =>
Type -> m (Kind ())
compileKind (Type -> m (Kind ())) -> Type -> m (Kind ())
forall a b. (a -> b) -> a -> b
$ Var -> Type
GHC.tyVarKind Var
v
  TyName
t' <- Name -> m TyName
forall (m :: * -> *). MonadQuote m => Name -> m TyName
compileTyNameFresh (Name -> m TyName) -> Name -> m TyName
forall a b. (a -> b) -> a -> b
$ Var -> Name
forall a. NamedThing a => a -> Name
GHC.getName Var
v
  PLCTyVar -> m PLCTyVar
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PLCTyVar -> m PLCTyVar) -> PLCTyVar -> m PLCTyVar
forall a b. (a -> b) -> a -> b
$ Ann -> TyName -> Kind Ann -> PLCTyVar
forall tyname ann.
ann -> tyname -> Kind ann -> TyVarDecl tyname ann
PLC.TyVarDecl Ann
annMayInline TyName
t' (Kind ()
k' Kind () -> Ann -> Kind Ann
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ann
annMayInline)

compileTcTyVarFresh :: (Compiling uni fun m ann) => GHC.TyCon -> m PLCTyVar
compileTcTyVarFresh :: forall (uni :: * -> *) fun (m :: * -> *) ann.
Compiling uni fun m ann =>
TyCon -> m PLCTyVar
compileTcTyVarFresh TyCon
tc = do
  Kind ()
k' <- Type -> m (Kind ())
forall (uni :: * -> *) fun (m :: * -> *) ann.
Compiling uni fun m ann =>
Type -> m (Kind ())
compileKind (Type -> m (Kind ())) -> Type -> m (Kind ())
forall a b. (a -> b) -> a -> b
$ TyCon -> Type
GHC.tyConKind TyCon
tc
  TyName
t' <- Name -> m TyName
forall (m :: * -> *). MonadQuote m => Name -> m TyName
compileTyNameFresh (Name -> m TyName) -> Name -> m TyName
forall a b. (a -> b) -> a -> b
$ TyCon -> Name
forall a. NamedThing a => a -> Name
GHC.getName TyCon
tc
  PLCTyVar -> m PLCTyVar
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PLCTyVar -> m PLCTyVar) -> PLCTyVar -> m PLCTyVar
forall a b. (a -> b) -> a -> b
$ Ann -> TyName -> Kind Ann -> PLCTyVar
forall tyname ann.
ann -> tyname -> Kind ann -> TyVarDecl tyname ann
PLC.TyVarDecl Ann
annMayInline TyName
t' (Kind ()
k' Kind () -> Ann -> Kind Ann
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ann
annMayInline)

pushName :: GHC.Name -> PLCVar uni -> Scope uni -> Scope uni
pushName :: forall (uni :: * -> *).
Name -> PLCVar uni -> Scope uni -> Scope uni
pushName Name
ghcName PLCVar uni
n (Scope Map Name (PLCVar uni)
ns Map Name PLCTyVar
tyns) = Map Name (PLCVar uni) -> Map Name PLCTyVar -> Scope uni
forall (uni :: * -> *).
Map Name (PLCVar uni) -> Map Name PLCTyVar -> Scope uni
Scope (Name
-> PLCVar uni -> Map Name (PLCVar uni) -> Map Name (PLCVar uni)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name
ghcName PLCVar uni
n Map Name (PLCVar uni)
ns) Map Name PLCTyVar
tyns

pushNames :: [(GHC.Name, PLCVar uni)] -> Scope uni -> Scope uni
pushNames :: forall (uni :: * -> *).
[(Name, PLCVar uni)] -> Scope uni -> Scope uni
pushNames [(Name, PLCVar uni)]
mappings Scope uni
scope = (Scope uni -> (Name, PLCVar uni) -> Scope uni)
-> Scope uni -> [(Name, PLCVar uni)] -> Scope uni
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Scope uni
acc (Name
n, PLCVar uni
v) -> Name -> PLCVar uni -> Scope uni -> Scope uni
forall (uni :: * -> *).
Name -> PLCVar uni -> Scope uni -> Scope uni
pushName Name
n PLCVar uni
v Scope uni
acc) Scope uni
scope [(Name, PLCVar uni)]
mappings

pushTyName :: GHC.Name -> PLCTyVar -> Scope uni -> Scope uni
pushTyName :: forall (uni :: * -> *). Name -> PLCTyVar -> Scope uni -> Scope uni
pushTyName Name
ghcName PLCTyVar
n (Scope Map Name (PLCVar uni)
ns Map Name PLCTyVar
tyns) = Map Name (PLCVar uni) -> Map Name PLCTyVar -> Scope uni
forall (uni :: * -> *).
Map Name (PLCVar uni) -> Map Name PLCTyVar -> Scope uni
Scope Map Name (PLCVar uni)
ns (Name -> PLCTyVar -> Map Name PLCTyVar -> Map Name PLCTyVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name
ghcName PLCTyVar
n Map Name PLCTyVar
tyns)

pushTyNames :: [(GHC.Name, PLCTyVar)] -> Scope uni -> Scope uni
pushTyNames :: forall (uni :: * -> *).
[(Name, PLCTyVar)] -> Scope uni -> Scope uni
pushTyNames [(Name, PLCTyVar)]
mappings Scope uni
scope = (Scope uni -> (Name, PLCTyVar) -> Scope uni)
-> Scope uni -> [(Name, PLCTyVar)] -> Scope uni
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Scope uni
acc (Name
n, PLCTyVar
v) -> Name -> PLCTyVar -> Scope uni -> Scope uni
forall (uni :: * -> *). Name -> PLCTyVar -> Scope uni -> Scope uni
pushTyName Name
n PLCTyVar
v Scope uni
acc) Scope uni
scope [(Name, PLCTyVar)]
mappings