Add name mangling pass to make VarIds globally unique (#120)

* Add name mangling pass to make VarIds globally unique

* Use slightly more esoteric mangling to avoid clashes

* Actually update the varNames map

* Make LambdaLift use manglenames

* Fix Lambda lift pass

* Fix things after different binder structure

* Still broken, waiting for types to appear in binders

* Ok ok ok ok

* Get it compiling

* Passtests

* Some small fixes

* Use symbol table to generate fresh names in InsertRefCounting

* Use symbol table to generate fresh names in LambdaLifting

* Use symbol table to generate fresh names in DesugarPattern

Author: j-hui

src/IR.hs

@@ -21,6 +21,7 @@ import IR.ExternToCall (externToCall)
 import IR.InsertRefCounting (insertRefCounting)
 import IR.LambdaLift (liftProgramLambdas)
 import IR.LowerAst (lowerProgram)
+import IR.MangleNames (mangleProgram)
 import IR.OptimizePar (optimizePar)
 import IR.Pattern (checkAnomaly)
 import IR.SegmentLets (segmentLets)
@@ -44,6 +45,7 @@ data Mode
   | DumpIRConstraints
   | DumpIRTyped
   | DumpIRTypedUgly
+  | DumpIRMangled
   | DumpIRInlined
   | DumpIRTypedShow
   | DumpIRLifted
@@ -88,6 +90,11 @@ options =
       ["dump-ir-typed-ugly"]
       (NoArg $ setMode DumpIRTypedShow)
       "Ugly-Print the fully-typed IR after type inference"
+  , Option
+      ""
+      ["dump-ir-mangled"]
+      (NoArg $ setMode DumpIRMangled)
+      "Print the IR after mangling"
   , Option
       ""
       ["dump-ir-lifted"]
@@ -139,6 +146,8 @@ anomalycheck p = do
 -- | IR transformations to prepare for codegen.
 transform :: Options -> I.Program I.Type -> Pass (I.Program I.Type)
 transform opt p = do
+  p <- mangleProgram p
+  when (mode opt == DumpIRMangled) $ dump p
   p <- desugarPattern p
   p <- instProgram p
   p <- segmentLets p

src/IR/Constraint/Constrain.hs

@@ -20,7 +20,7 @@ sprinkleVariables
   :: I.Program Can.Annotations -> TC (I.Program (Can.Annotations, Variable))
 sprinkleVariables prog = do
   sprinkledDefs <- mapM sprinkleDef (I.programDefs prog)
-  return prog{I.programDefs = sprinkledDefs}
+  return prog{I.programDefs = sprinkledDefs, I.symTable = I.uninitializedSymTable}
  where
   sprinkleDef (name, expr) = do
     name' <- mapM sprinkle name
@@ -35,4 +35,4 @@ discardAnnotations
   :: I.Program (Can.Annotations, Variable) -> I.Program Variable
 discardAnnotations sprinkledProg =
   let discardedDefs = map (bimap (fmap snd) (fmap snd)) (I.programDefs sprinkledProg)
-   in sprinkledProg{I.programDefs = discardedDefs}
+   in sprinkledProg{I.programDefs = discardedDefs, I.symTable = I.uninitializedSymTable}

src/IR/Constraint/Elaborate.hs

@@ -12,4 +12,4 @@ run pVar = do
   exprs' <- mapM (mapM Type.toCanType) exprs
   let reattachBinders (I.BindVar v _) expr = (I.BindVar v $ I.extract expr, expr)
       reattachBinders (I.BindAnon _) expr = (I.BindAnon $ I.extract expr, expr)
-  return $ pVar{I.programDefs = zipWith reattachBinders names exprs'}
+  return $ pVar{I.programDefs = zipWith reattachBinders names exprs', I.symTable = I.uninitializedSymTable }

src/IR/DesugarPattern.hs

@@ -13,21 +13,19 @@ import Common.Compiler as Compiler (
   Pass,
   fromString,
  )
-import Common.Identifiers (Identifier (..))
 import qualified IR.IR as I
+import qualified IR.MangleNames as I
 
 import Common.Pretty (Pretty (..))
-import Control.Monad (forM, replicateM)
+import Control.Monad (forM)
 import Control.Monad.State.Lazy (
   MonadState,
   StateT (..),
-  evalStateT,
   gets,
   modify,
  )
 import Data.Bifunctor (Bifunctor (..))
 import Data.Foldable (foldrM)
-import Data.Functor ((<&>))
 import qualified Data.Map as M
 import qualified Data.Set as S
 
@@ -39,7 +37,6 @@ data CInfo = CInfo
   { cName :: I.DConId
   , cType :: I.TConId
   , argsType :: [I.Type]
-  , cArity :: Int
   }
   deriving (Eq, Show)
 
@@ -54,7 +51,7 @@ data TInfo = TInfo
 data DesugarCtx = DesugarCtx
   { typeMap :: M.Map I.TConId TInfo
   , consMap :: M.Map I.DConId CInfo
-  , anonCount :: Int
+  , symTable :: I.SymTable I.Type
   }
 
 
@@ -67,24 +64,18 @@ newtype DesugarFn a = DesugarFn (StateT DesugarCtx Compiler.Pass a)
   deriving (MonadState DesugarCtx) via (StateT DesugarCtx Compiler.Pass)
 
 
-runDesugarFn :: DesugarFn a -> DesugarCtx -> Compiler.Pass a
-runDesugarFn (DesugarFn m) = evalStateT m
+unDesugarFn :: DesugarFn a -> StateT DesugarCtx Compiler.Pass a
+unDesugarFn (DesugarFn a) = a
 
 
-freshVar :: DesugarFn Identifier
-freshVar = do
-  currCount <- gets anonCount
-  modify $ \ctx -> ctx{anonCount = anonCount ctx + 1}
-  return $ fromString ("__pat_anon" ++ show currCount)
-
-
-buildCtx :: [(I.TConId, I.TypeDef)] -> DesugarCtx
-buildCtx tds =
-  DesugarCtx
-    { anonCount = 0
-    , typeMap = buildTypeMap tds
-    , consMap = buildConsMap tds
-    }
+freshVar :: I.Type -> DesugarFn I.VarId
+freshVar t = do
+  syms <- gets symTable
+  let origin = "__anonymous_pattern"
+      name = I.pickId syms origin
+      syms' = M.insert name I.SymInfo{I.symOrigin = origin, I.symType = t} syms
+  modify $ \ctx -> ctx{symTable = syms'}
+  return name
 
 
 -- TODO: this should be a pattern synonym
@@ -93,12 +84,12 @@ unreachableExpr = I.Exception $ I.ExceptDefault $ I.LitIntegral 0
 
 
 desugarPattern :: I.Program I.Type -> Compiler.Pass (I.Program I.Type)
-desugarPattern p@I.Program{I.programDefs = defs, I.typeDefs = tds} =
-  (`runDesugarFn` buildCtx tds) $ do
-    defs' <- mapM desugarExprsDefs defs
-    return $ p{I.programDefs = defs'}
+desugarPattern p@I.Program{I.programDefs = defs, I.typeDefs = tds, I.symTable = syms} = do
+  (defs', symTable -> syms') <- runStateT (unDesugarFn $ mapM desugarExprsDefs defs) initCtx
+  return $ p{I.programDefs = defs', I.symTable = syms'}
  where
-  desugarExprsDefs (vs, es) = desugarExpr es <&> (vs,)
+  desugarExprsDefs (vs, es) = (vs,) <$> desugarExpr es
+  initCtx = DesugarCtx{typeMap = buildTypeMap tds, consMap = buildConsMap tds, symTable = syms}
 
 
 desugarExpr :: I.Expr I.Type -> DesugarFn (I.Expr I.Type)
@@ -116,19 +107,19 @@ desugarExpr (I.Match e arms t) = do
     I.Var _ _ -> desugarMatch [e] eqns (unreachableExpr t) -- TODO: add let alias
     _ -> do
       -- Bind scrutinee to a variable before threading it through desugarMatch
-      var <- I.VarId <$> freshVar
       let et = I.extract e
+      var <- freshVar et
       I.Let [(I.BindVar var et, e)]
         <$> desugarMatch [I.Var var et] eqns (unreachableExpr t)
         <*> pure et
 desugarExpr e = return e
 
 
-desugarMatch
-  :: [I.Expr I.Type]
-  -> [Equation]
-  -> I.Expr I.Type -- Default expression the 'I.Match' should return
-  -> DesugarFn (I.Expr I.Type)
+desugarMatch ::
+  [I.Expr I.Type] ->
+  [Equation] ->
+  I.Expr I.Type -> -- Default expression the 'I.Match' should return
+  DesugarFn (I.Expr I.Type)
 desugarMatch [] [] def = return def
 desugarMatch [] (([], e) : _) _ = return e
 desugarMatch [] _ _ = error "can't happen 1"
@@ -190,27 +181,28 @@ desugarMatchCons (u : us) qs@(q : _) def = do
   getCon ((I.AltData dcon _ _) : _, _) = dcon
   getCon _ = error "can't happen 5"
   getTyp ((I.AltData _ _ t) : _, _) = t
-  getTyp _ = error "no no no -- Simon Peyton Jones"
-  makeBinder vid t = I.AltBinder $ I.BindVar vid t
+  getTyp _ = error "no no no 555"
 
   sameConsAs c = filter ((== c) . getCon) qs
 
   desugarArm :: I.DConId -> I.Type -> [Equation] -> DesugarFn (I.Alt I.Type, I.Expr I.Type)
   desugarArm dcon dconTyp qs' = do
-    k <- getArity dcon
-    newIds <- replicateM k freshVar
-    cinfo <- getCInfo dcon
-    let newVars = I.VarId <$> newIds
-        argsTyps = argsType cinfo
-        bs' = zipWith makeBinder newVars argsTyps
-        us' = zipWith I.Var newVars argsTyps
-        qs'' = [(as' ++ as, e) | ((I.AltData _ as' _) : as, e) <- qs']
+    argsTyps <- argsType <$> getCInfo dcon
+
+    (unzip -> (bs', us')) <- forM argsTyps $ \argTyp -> do
+      name <- freshVar argTyp
+      return (I.AltBinder $ I.BindVar name argTyp, I.Var name argTyp)
+
     body <-
       if null qs'
-        then -- We're done desugaring this equation, just use default body
+        then do
+          -- We're done desugaring this equation, just use default body
           return def
-        else -- Recursively generate body from remaining equations
+        else do
+          -- Recursively generate body from remaining equations
+          let qs'' = [(as' ++ as, e) | ((I.AltData _ as' _) : as, e) <- qs']
           desugarMatch (us' ++ us) qs'' def
+
     return (I.AltData dcon bs' dconTyp, body)
 desugarMatchCons _ _ _ = error "can't happen 7"
 
@@ -232,12 +224,6 @@ getConstructors dcon = do
   return $ tCSet t
 
 
-getArity :: I.DConId -> DesugarFn Int
-getArity dcon = do
-  c <- getCInfo dcon
-  return $ cArity c
-
-
 getCInfo :: I.DConId -> DesugarFn CInfo
 getCInfo dcon = gets (M.lookup dcon . consMap) >>= maybe (desugarError dcon) return
 
@@ -269,5 +255,4 @@ buildConsMap = foldr (build . second I.variants) M.empty
         { cName = dcon
         , cType = tcon
         , argsType = variantTypes typs
-        , cArity = length $ variantTypes typs
         }

src/IR/IR.hs

@@ -8,6 +8,7 @@
 -- | Sslang's intermediate representation and its associated helpers.
 module IR.IR (
   Program (..),
+  SymInfo (..),
   TypeDef (..),
   TypeVariant (..),
   Binder (..),
@@ -37,6 +38,8 @@ module IR.IR (
   pattern BindAnon,
   binderToVar,
   Carrier,
+  uninitializedSymTable,
+  SymTable,
 ) where
 
 import Common.Identifiers (
@@ -52,6 +55,7 @@ import Data.Data (
   Typeable,
  )
 
+import qualified Data.Map as M
 import Data.Maybe (
   catMaybes,
   mapMaybe,
@@ -67,16 +71,31 @@ import IR.Types.Type (
  )
 
 
-{- | Top-level compilation unit.
-
-@t@ is the type system in use, e.g., "IR.Types.Flat"
--}
+-- | Top-level compilation unit, parameterized by the type system.
 data Program t = Program
   { programEntry :: VarId
   , cDefs :: String
   , externDecls :: [(VarId, Type)]
   , programDefs :: [(Binder t, Expr t)]
   , typeDefs :: [(TConId, TypeDef)]
+  , symTable :: M.Map VarId (SymInfo t)
+  }
+  deriving (Eq, Show, Typeable, Data, Functor, Foldable, Traversable)
+
+
+{- | Contains information about all (globally unique) variable names.
+
+Populated by name mangling pass.
+-}
+type SymTable t = M.Map VarId (SymInfo t)
+
+
+-- | Information stored in global symbol table.
+data SymInfo t = SymInfo
+  { -- | Original name of symbol
+    symOrigin :: VarId
+  , -- | Type information of symbol
+    symType :: t
   }
   deriving (Eq, Show, Typeable, Data, Functor, Foldable, Traversable)
 
@@ -409,6 +428,10 @@ binderToVar (BindVar v _) = Just v
 binderToVar _ = Nothing
 
 
+uninitializedSymTable :: M.Map VarId (SymInfo t)
+uninitializedSymTable = M.empty
+
+
 instance HasFreeVars (Expr t) VarId where
   freeVars (Var v _) = S.singleton v
   freeVars Data{} = S.empty