new grammar and adapted renamer

src/Renamer/Renamer.hs

@@ -1,56 +1,101 @@
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PatternSynonyms #-}
 
-module Renamer.Renamer where
+module Renamer.Renamer (rename) where
 
 import Auxiliary (mapAccumM)
+import Control.Applicative (Applicative (liftA2))
+import Control.Monad.Except (ExceptT, runExceptT)
+import Control.Monad.Identity (Identity, runIdentity)
 import Control.Monad.State (
     MonadState,
-    State,
-    evalState,
+    StateT,
+    evalStateT,
     gets,
     modify,
  )
-import Data.List (foldl')
+import Data.Function (on)
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe)
 import Data.Tuple.Extra (dupe)
-import Debug.Trace (trace)
 import Grammar.Abs
 
 -- | Rename all variables and local binds
-rename :: Program -> Program
-rename (Program bs) = Program $ evalState (runRn $ mapM (renameSc initNames) bs) 0
+rename :: Program -> Either String Program
+rename (Program defs) = Program <$> renameDefs defs
+
+renameDefs :: [Def] -> Either String [Def]
+renameDefs defs = runIdentity $ runExceptT $ evalStateT (runRn $ mapM renameDef defs) initCxt
   where
-    -- initNames = Map.fromList $ map (\(Bind name _ _ _ _) -> dupe name) bs
-    initNames = Map.fromList $ foldl' saveIfBind [] bs
-    saveIfBind acc (DBind (Bind name _ _ _ _)) = dupe name : acc
-    saveIfBind acc _ = acc
-    renameSc :: Names -> Def -> Rn Def
-    renameSc old_names (DBind (Bind name t _ parms rhs)) = do
-        (new_names, parms') <- newNames old_names parms
-        rhs' <- snd <$> renameExp new_names rhs
-        pure . DBind $ Bind name t name parms' rhs'
-    renameSc _ def = pure def
+    initNames = Map.fromList [dupe name | DBind (Bind name _ _) <- defs]
+
+    renameDef :: Def -> Rn Def
+    renameDef = \case
+        DSig (Sig name typ) -> DSig . Sig name <$> renameTVars typ
+        DBind (Bind name vars rhs) -> do
+            (new_names, vars') <- newNames initNames vars
+            rhs' <- snd <$> renameExp new_names rhs
+            pure . DBind $ Bind name vars' rhs'
+        DData (Data (Indexed cname types) constrs) -> do
+            tvars' <- mapM nextNameTVar tvars
+            let tvars_lt = zip tvars tvars'
+                typ' = map (substituteTVar tvars_lt) types
+                constrs' = map (renameConstr tvars_lt) constrs
+            pure . DData $ Data (Indexed cname typ') constrs'
+          where
+            tvars = concatMap (collectTVars []) types
+            collectTVars tvars = \case
+                TAll tvar t -> collectTVars (tvar : tvars) t
+                TIndexed _ -> tvars
+                -- Should be monad error
+                TVar v -> [v]
+                _ -> error ("Bad data type definition: " ++ show types)
+
+    renameConstr :: [(TVar, TVar)] -> Constructor -> Constructor
+    renameConstr new_types (Constructor name typ) =
+        Constructor name $ substituteTVar new_types typ
+
+substituteTVar :: [(TVar, TVar)] -> Type -> Type
+substituteTVar new_names typ = case typ of
+    TLit _ -> typ
+    TVar tvar
+        | Just tvar' <- lookup tvar new_names ->
+            TVar tvar'
+        | otherwise ->
+            typ
+    TFun t1 t2 -> on TFun substitute' t1 t2
+    TAll tvar t
+        | Just tvar' <- lookup tvar new_names ->
+            TAll tvar' $ substitute' t
+        | otherwise ->
+            TAll tvar $ substitute' t
+    TIndexed (Indexed name typs) -> TIndexed . Indexed name $ map substitute' typs
+    _ -> error ("Impossible " ++ show typ)
+  where
+    substitute' = substituteTVar new_names
+
+initCxt :: Cxt
+initCxt = Cxt 0 0
+
+data Cxt = Cxt
+    { var_counter :: Int
+    , tvar_counter :: Int
+    }
 
 -- | Rename monad. State holds the number of renamed names.
-newtype Rn a = Rn {runRn :: State Int a}
-    deriving (Functor, Applicative, Monad, MonadState Int)
+newtype Rn a = Rn {runRn :: StateT Cxt (ExceptT String Identity) a}
+    deriving (Functor, Applicative, Monad, MonadState Cxt)
 
 -- | Maps old to new name
-type Names = Map Ident Ident
-
-renameLocalBind :: Names -> Bind -> Rn (Names, Bind)
-renameLocalBind old_names (Bind name t _ parms rhs) = do
-    (new_names, name') <- newName old_names name
-    (new_names', parms') <- newNames new_names parms
-    (new_names'', rhs') <- renameExp new_names' rhs
-    pure (new_names'', Bind name' t name' parms' rhs')
+type Names = Map LIdent LIdent
 
 renameExp :: Names -> Exp -> Rn (Names, Exp)
 renameExp old_names = \case
     EId n -> pure (old_names, EId . fromMaybe n $ Map.lookup n old_names)
-    ELit (LInt i1) -> pure (old_names, ELit (LInt i1))
+    ELit lit -> pure (old_names, ELit lit)
     EApp e1 e2 -> do
         (env1, e1') <- renameExp old_names e1
         (env2, e2') <- renameExp old_names e2
@@ -59,17 +104,21 @@ renameExp old_names = \case
         (env1, e1') <- renameExp old_names e1
         (env2, e2') <- renameExp old_names e2
         pure (Map.union env1 env2, EAdd e1' e2')
-    ELet i e1 e2 -> do
-        (new_names, e1') <- renameExp old_names e1
-        (new_names', e2') <- renameExp new_names e2
-        pure (new_names', ELet i e1' e2')
+
+    -- TODO fix shadowing
+    ELet name rhs e -> do
+        (new_names, name') <- newName old_names name
+        (new_names', rhs') <- renameExp new_names rhs
+        (new_names'', e') <- renameExp new_names' e
+        pure (new_names'', ELet name' rhs' e')
     EAbs par e -> do
         (new_names, par') <- newName old_names par
         (new_names', e') <- renameExp new_names e
         pure (new_names', EAbs par' e')
     EAnn e t -> do
         (new_names, e') <- renameExp old_names e
-        pure (new_names, EAnn e' t)
+        t' <- renameTVars t
+        pure (new_names, EAnn e' t')
     ECase e injs -> do
         (new_names, e') <- renameExp old_names e
         (new_names', injs') <- renameInjs new_names injs
@@ -88,21 +137,58 @@ renameInj ns (Inj init e) = do
 
 renameInit :: Names -> Init -> Rn (Names, Init)
 renameInit ns i = case i of
-    InitConstr cs vars -> do
+    InitConstructor cs vars -> do
         (ns_new, vars') <- newNames ns vars
-        return (ns_new, InitConstr cs vars')
+        return (ns_new, InitConstructor cs vars')
     rest -> return (ns, rest)
 
+renameTVars :: Type -> Rn Type
+renameTVars typ = case typ of
+    TAll tvar t -> do
+        tvar' <- nextNameTVar tvar
+        t' <- renameTVars $ substitute tvar tvar' t
+        pure $ TAll tvar' t'
+    TFun t1 t2 -> liftA2 TFun (renameTVars t1) (renameTVars t2)
+    _ -> pure typ
+
+substitute ::
+    TVar -> -- α
+    TVar -> -- α_n
+    Type -> -- A
+    Type -- [α_n/α]A
+substitute tvar1 tvar2 typ = case typ of
+    TLit _ -> typ
+    TVar tvar'
+        | tvar' == tvar1 -> TVar tvar2
+        | otherwise -> typ
+    TFun t1 t2 -> on TFun substitute' t1 t2
+    TAll tvar t -> TAll tvar $ substitute' t
+    TIndexed (Indexed name typs) -> TIndexed . Indexed name $ map substitute' typs
+    _ -> error "Impossible"
+  where
+    substitute' = substitute tvar1 tvar2
+
 -- | Create a new name and add it to name environment.
-newName :: Names -> Ident -> Rn (Names, Ident)
+newName :: Names -> LIdent -> Rn (Names, LIdent)
 newName env old_name = do
     new_name <- makeName old_name
     pure (Map.insert old_name new_name env, new_name)
 
 -- | Create multiple names and add them to the name environment
-newNames :: Names -> [Ident] -> Rn (Names, [Ident])
+newNames :: Names -> [LIdent] -> Rn (Names, [LIdent])
 newNames = mapAccumM newName
 
 -- | Annotate name with number and increment the number @prefix ⇒ prefix_number@.
-makeName :: Ident -> Rn Ident
-makeName (Ident prefix) = gets (\i -> Ident $ prefix ++ "_" ++ show i) <* modify succ
+makeName :: LIdent -> Rn LIdent
+makeName (LIdent prefix) = do
+    i <- gets var_counter
+    let name = LIdent $ prefix ++ "_" ++ show i
+    modify $ \cxt -> cxt{var_counter = succ cxt.var_counter}
+    pure name
+
+nextNameTVar :: TVar -> Rn TVar
+nextNameTVar (MkTVar (LIdent s)) = do
+    i <- gets tvar_counter
+    let tvar = MkTVar . LIdent $ s ++ "_" ++ show i
+    modify $ \cxt -> cxt{tvar_counter = succ cxt.tvar_counter}
+    pure tvar