{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module      : Jikka.Core.Convert.Alpha
-- Description : does alpha-conversion. / alpha 変換をします。
-- Copyright   : (c) Kimiyuki Onaka, 2020
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Convert.Alpha
  ( run,
    runProgram,
    runToplevelExpr,
    runExpr,
  )
where

import Control.Monad.State.Strict
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint

type UsedVars = [VarName]

type RenameMapping = [(VarName, VarName)]

rename :: (MonadState UsedVars m, MonadAlpha m) => VarName -> m VarName
rename :: VarName -> m VarName
rename VarName
x = do
  [VarName]
used <- m [VarName]
forall s (m :: * -> *). MonadState s m => m s
get
  VarName
y <-
    if VarName
x VarName -> [VarName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VarName]
used
      then VarName -> m VarName
forall (m :: * -> *) a. Monad m => a -> m a
return VarName
x
      else do
        let base :: [Char]
base = (Char -> Bool) -> [Char] -> [Char]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'$') (VarName -> [Char]
formatVarName VarName
x)
        OccName -> NameFlavour -> VarName
VarName ([Char] -> OccName
forall a. a -> Maybe a
Just [Char]
base) (NameFlavour -> VarName) -> (Int -> NameFlavour) -> Int -> VarName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> NameFlavour
forall a. a -> Maybe a
Just (Int -> VarName) -> m Int -> m VarName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Int
forall (m :: * -> *). MonadAlpha m => m Int
nextCounter
  [VarName] -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([VarName] -> m ()) -> [VarName] -> m ()
forall a b. (a -> b) -> a -> b
$ VarName
y VarName -> [VarName] -> [VarName]
forall a. a -> [a] -> [a]
: [VarName]
used
  VarName -> m VarName
forall (m :: * -> *) a. Monad m => a -> m a
return VarName
y

runExpr' :: (MonadState UsedVars m, MonadAlpha m, MonadError Error m) => RenameMapping -> Expr -> m Expr
runExpr' :: RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env = \case
  Var VarName
x -> case VarName -> RenameMapping -> Maybe VarName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup VarName
x RenameMapping
env of
    Maybe VarName
Nothing -> [Char] -> m Expr
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a
throwInternalError ([Char] -> m Expr) -> [Char] -> m Expr
forall a b. (a -> b) -> a -> b
$ [Char]
"undefined variable: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VarName -> [Char]
formatVarName VarName
x
    Just VarName
y -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
y
  Lit Literal
lit -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
  App Expr
f Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
f m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e
  Lam VarName
x Type
t Expr
body -> do
    VarName
y <- VarName -> m VarName
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m) =>
VarName -> m VarName
rename VarName
x
    Expr
body <- RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' ((VarName
x, VarName
y) (VarName, VarName) -> RenameMapping -> RenameMapping
forall a. a -> [a] -> [a]
: RenameMapping
env) Expr
body
    Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> Expr -> Expr
Lam VarName
y Type
t Expr
body
  Let VarName
x Type
t Expr
e1 Expr
e2 -> do
    Expr
e1 <- RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e1
    VarName
y <- VarName -> m VarName
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m) =>
VarName -> m VarName
rename VarName
x
    Expr
e2 <- RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' ((VarName
x, VarName
y) (VarName, VarName) -> RenameMapping -> RenameMapping
forall a. a -> [a] -> [a]
: RenameMapping
env) Expr
e2
    Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> Expr -> Expr -> Expr
Let VarName
y Type
t Expr
e1 Expr
e2
  Assert Expr
e1 Expr
e2 -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e2

runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
runExpr :: [(VarName, Type)] -> Expr -> m Expr
runExpr [(VarName, Type)]
env Expr
e = [Char] -> m Expr -> m Expr
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a -> m a
wrapError' [Char]
"Jikka.Core.Convert.Alpha.runExpr" (m Expr -> m Expr) -> m Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ do
  StateT [VarName] m Expr -> [VarName] -> m Expr
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (RenameMapping -> Expr -> StateT [VarName] m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' (((VarName, Type) -> (VarName, VarName))
-> [(VarName, Type)] -> RenameMapping
forall a b. (a -> b) -> [a] -> [b]
map (\(VarName
x, Type
_) -> (VarName
x, VarName
x)) [(VarName, Type)]
env) Expr
e) (((VarName, Type) -> VarName) -> [(VarName, Type)] -> [VarName]
forall a b. (a -> b) -> [a] -> [b]
map (VarName, Type) -> VarName
forall a b. (a, b) -> a
fst [(VarName, Type)]
env)

runToplevelExpr' :: (MonadState UsedVars m, MonadAlpha m, MonadError Error m) => RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' :: RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' RenameMapping
env = \case
  ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> m Expr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e
  ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont -> do
    VarName
y <- VarName -> m VarName
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m) =>
VarName -> m VarName
rename VarName
x
    Expr
e <- RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e
    ToplevelExpr
cont <- RenameMapping -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' ((VarName
x, VarName
y) (VarName, VarName) -> RenameMapping -> RenameMapping
forall a. a -> [a] -> [a]
: RenameMapping
env) ToplevelExpr
cont
    ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
y Type
t Expr
e ToplevelExpr
cont
  ToplevelLetRec VarName
f [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont -> do
    VarName
g <- VarName -> m VarName
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m) =>
VarName -> m VarName
rename VarName
f
    [(VarName, VarName, Type)]
args <- [(VarName, Type)]
-> ((VarName, Type) -> m (VarName, VarName, Type))
-> m [(VarName, VarName, Type)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(VarName, Type)]
args (((VarName, Type) -> m (VarName, VarName, Type))
 -> m [(VarName, VarName, Type)])
-> ((VarName, Type) -> m (VarName, VarName, Type))
-> m [(VarName, VarName, Type)]
forall a b. (a -> b) -> a -> b
$ \(VarName
x, Type
t) -> do
      VarName
y <- VarName -> m VarName
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m) =>
VarName -> m VarName
rename VarName
x
      (VarName, VarName, Type) -> m (VarName, VarName, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName
x, VarName
y, Type
t)
    let args1 :: RenameMapping
args1 = ((VarName, VarName, Type) -> (VarName, VarName))
-> [(VarName, VarName, Type)] -> RenameMapping
forall a b. (a -> b) -> [a] -> [b]
map (\(VarName
x, VarName
y, Type
_) -> (VarName
x, VarName
y)) [(VarName, VarName, Type)]
args
    let args2 :: [(VarName, Type)]
args2 = ((VarName, VarName, Type) -> (VarName, Type))
-> [(VarName, VarName, Type)] -> [(VarName, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(VarName
_, VarName
y, Type
t) -> (VarName
y, Type
t)) [(VarName, VarName, Type)]
args
    Expr
body <- RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' (RenameMapping
args1 RenameMapping -> RenameMapping -> RenameMapping
forall a. [a] -> [a] -> [a]
++ (VarName
f, VarName
g) (VarName, VarName) -> RenameMapping -> RenameMapping
forall a. a -> [a] -> [a]
: RenameMapping
env) Expr
body
    ToplevelExpr
cont <- RenameMapping -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' ((VarName
f, VarName
g) (VarName, VarName) -> RenameMapping -> RenameMapping
forall a. a -> [a] -> [a]
: RenameMapping
env) ToplevelExpr
cont
    ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
g [(VarName, Type)]
args2 Type
ret Expr
body ToplevelExpr
cont
  ToplevelAssert Expr
e1 ToplevelExpr
e2 -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RenameMapping -> Expr -> m Expr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> Expr -> m Expr
runExpr' RenameMapping
env Expr
e1 m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RenameMapping -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' RenameMapping
env ToplevelExpr
e2

runToplevelExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr :: [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr [(VarName, Type)]
env ToplevelExpr
e = [Char] -> m ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a -> m a
wrapError' [Char]
"Jikka.Core.Convert.Alpha.runToplevelExpr" (m ToplevelExpr -> m ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ do
  StateT [VarName] m ToplevelExpr -> [VarName] -> m ToplevelExpr
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (RenameMapping -> ToplevelExpr -> StateT [VarName] m ToplevelExpr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' (((VarName, Type) -> (VarName, VarName))
-> [(VarName, Type)] -> RenameMapping
forall a b. (a -> b) -> [a] -> [b]
map (\(VarName
x, Type
_) -> (VarName
x, VarName
x)) [(VarName, Type)]
env) ToplevelExpr
e) (((VarName, Type) -> VarName) -> [(VarName, Type)] -> [VarName]
forall a b. (a -> b) -> [a] -> [b]
map (VarName, Type) -> VarName
forall a b. (a, b) -> a
fst [(VarName, Type)]
env)

runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
runProgram :: ToplevelExpr -> m ToplevelExpr
runProgram ToplevelExpr
prog = [Char] -> m ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. MonadError Error m => [Char] -> m a -> m a
wrapError' [Char]
"Jikka.Core.Convert.Alpha.runProgram" (m ToplevelExpr -> m ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ do
  ToplevelExpr
prog <- StateT [VarName] m ToplevelExpr -> [VarName] -> m ToplevelExpr
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (RenameMapping -> ToplevelExpr -> StateT [VarName] m ToplevelExpr
forall (m :: * -> *).
(MonadState [VarName] m, MonadAlpha m, MonadError Error m) =>
RenameMapping -> ToplevelExpr -> m ToplevelExpr
runToplevelExpr' [] ToplevelExpr
prog) []
  m () -> m ()
forall (m :: * -> *) a. MonadError Error m => m a -> m a
postcondition (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
ensureAlphaConverted ToplevelExpr
prog
  ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
prog

-- | `run` renames variables in exprs to avoid name conflictions, even if the scopes of two variables are distinct.
--
-- == Examples
--
-- Before:
--
-- > let x = 0
-- > in y = x + x
-- > in x = x + y
-- > x + y
--
-- After:
--
-- > let x0 = 0
-- > in y1 = x0 + x0
-- > in x2 = x0 + y1
-- > x2 + y1
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run :: ToplevelExpr -> m ToplevelExpr
run = ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
(MonadAlpha m, MonadError Error m) =>
ToplevelExpr -> m ToplevelExpr
runProgram