From 77749172588c6d99b5dbbcc5838afbb4602f59ad Mon Sep 17 00:00:00 2001
From: Ben Price <ben@hackworthltd.com>
Date: Mon, 20 Jun 2022 15:19:28 +0100
Subject: [PATCH] refactor: common up uniquifying names

Note that this changes the results of 'uniquifyDefName': previously we
would have generated 'def_1', and now we generate 'def1' (without an
underscore).
---
 primer/src/Primer/Action.hs    | 20 ++++++++------------
 primer/src/Primer/Questions.hs | 20 ++++++++++++++++----
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/primer/src/Primer/Action.hs b/primer/src/Primer/Action.hs
index faf810032..acca1228d 100644
--- a/primer/src/Primer/Action.hs
+++ b/primer/src/Primer/Action.hs
@@ -27,6 +27,7 @@ import Data.Aeson (Value)
 import Data.Generics.Product (typed)
 import Data.List (findIndex)
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import Optics (set, (%), (?~))
 import Primer.Core (
@@ -87,14 +88,14 @@ import Primer.Core.Transform (renameLocalVar, renameTyVar, renameTyVarExpr)
 import Primer.Core.Utils (forgetTypeIDs, generateTypeIDs)
 import Primer.JSON
 import Primer.Module (Module, insertDef)
-import Primer.Name (Name, NameCounter, unName)
+import Primer.Name (Name, NameCounter, unName, unsafeMkName)
 import Primer.Name.Fresh (
   isFresh,
   isFreshTy,
   mkFreshName,
   mkFreshNameTy,
  )
-import Primer.Questions (Question)
+import Primer.Questions (Question, uniquify)
 import Primer.Refine (Inst (InstAPP, InstApp, InstUnconstrainedAPP), refine)
 import Primer.Typecheck (
   CheckEverythingRequest (CheckEverything, toCheck, trusted),
@@ -232,18 +233,13 @@ nameString = "n" <> T.singleton '\x200C' <> "ame"
 -- of the given name already exists in the program, this function will
 -- return the same name it's been given.
 uniquifyDefName :: C.ModuleName -> Text -> DefMap -> Text
-uniquifyDefName m name' defs =
-  if name' `notElem` avoid
-    then name'
-    else
-      let go i = if (name' <> "_" <> show i) `notElem` avoid then name' <> "_" <> show i else go (i + 1)
-       in go (1 :: Int)
+uniquifyDefName m name' defs = unName $ uniquify avoid $ unsafeMkName name'
   where
     f qn
-      | qualifiedModule qn == m = Just (unName $ baseName qn)
-      | otherwise = Nothing
-    avoid :: [Text]
-    avoid = mapMaybe f $ Map.keys defs
+      | qualifiedModule qn == m = Set.singleton $ baseName qn
+      | otherwise = mempty
+    avoid :: Set Name
+    avoid = foldMap f $ Map.keys defs
 
 type QualifiedText = (NonEmpty Text, Text)
 
diff --git a/primer/src/Primer/Questions.hs b/primer/src/Primer/Questions.hs
index 74b7b8fb7..141dbb43d 100644
--- a/primer/src/Primer/Questions.hs
+++ b/primer/src/Primer/Questions.hs
@@ -10,6 +10,7 @@ module Primer.Questions (
   ShadowedVarsTy (..), -- only exported for testing
   generateNameExpr,
   generateNameTy,
+  uniquify,
 ) where
 
 import Foreword
@@ -90,7 +91,7 @@ generateNameExpr ::
 -- in a term context (second Either is Left): we could be inserting a LAM.
 -- It doesn't make sense to ask for a term variable in a type context,
 -- but it also doesn't harm to support it.
-generateNameExpr tk z = uniquify <$> getAvoidSet z <*> baseNames tk
+generateNameExpr tk z = uniquifyMany <$> getAvoidSet z <*> baseNames tk
 
 generateNameTy ::
   MonadReader Cxt m =>
@@ -99,7 +100,7 @@ generateNameTy ::
   m [Name]
 -- It doesn't really make sense to ask for a term variable (Left) here, but
 -- it doesn't harm to support it
-generateNameTy tk z = uniquify <$> mkAvoidForFreshNameTy z <*> baseNames tk
+generateNameTy tk z = uniquifyMany <$> mkAvoidForFreshNameTy z <*> baseNames tk
 
 baseNames ::
   MonadReader Cxt m =>
@@ -125,10 +126,16 @@ getAvoidSet = \case
   Left ze -> mkAvoidForFreshName ze
   Right zt -> mkAvoidForFreshNameTypeZ zt
 
+-- | Adds a numeric suffix to a name to be distinct from a given set.
+-- (If the name is already distinct then return it unmodified.)
+uniquify :: Set.Set Name -> Name -> Name
+uniquify avoid = snd . uniquify' avoid
+
+-- A helper for uniquify and uniquifyMany
 -- We do not use Name.freshName as we don't want a global fresh counter
 -- (and we want to control the base name)
-uniquify :: Set.Set Name -> [Name] -> [Name]
-uniquify avoid ns = map snd $ sort $ map go ns
+uniquify' :: Set.Set Name -> Name -> (Integer, Name)
+uniquify' avoid = go
   where
     -- Replace use of `unsafeHead` here. See:
     -- https://github.com/hackworthltd/primer/issues/147
@@ -137,3 +144,8 @@ uniquify avoid ns = map snd $ sort $ map go ns
     f n = \case
       0 -> n
       i -> unsafeMkName $ unName n <> show i
+
+-- | Adds a numeric suffix to each name so they are distinct from the given set.
+-- Returns the thus-constructed names in order of their added suffix.
+uniquifyMany :: Set.Set Name -> [Name] -> [Name]
+uniquifyMany avoid ns = map snd $ sort $ uniquify' avoid <$> ns