From 3ee2d5c5308d37c712cdef43f9c15a07fe68644b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samuel=20G=C3=A9lineau?= <samuel.gelineau@well.co>
Date: Sun, 9 Jul 2023 20:20:59 -0400
Subject: [PATCH] lookupTyCon, lookupDataCon

I doubt the implementation I gave will work on older versions of GHC,
let's see where CI breaks.
---
 src/Hint/Base.hs | 22 ++++++++++++++++++++-
 src/Hint/GHC.hs  | 50 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/Hint/Base.hs b/src/Hint/Base.hs
index f1bafc7..ebd2ddf 100644
--- a/src/Hint/Base.hs
+++ b/src/Hint/Base.hs
@@ -9,7 +9,7 @@ module Hint.Base (
     ImportList(..), ModuleQualification(..), ModuleImport(..),
 
     ModuleName, PhantomModule(..),
-    findModule, moduleIsLoaded,
+    findModule, lookupTyCon, lookupDataCon, moduleIsLoaded,
     withDynFlags,
 
     ghcVersion,
@@ -175,6 +175,26 @@ findModule mn = mapGhcExceptions NotAllowed $
                     runGhc $ GHC.findModule mod_name Nothing
     where mod_name = GHC.mkModuleName mn
 
+lookupTyCon :: MonadInterpreter m => ModuleName -> String -> m GHC.TyCon
+lookupTyCon moduleName tyConName = do
+  mapGhcExceptions NotAllowed $ do
+    runGhc $ do
+      GHC.lookupTyCon moduleName tyConName >>= \case
+        Right tyCon -> do
+          pure tyCon
+        Left err -> do
+          throwM $ UnknownError err
+
+lookupDataCon :: MonadInterpreter m => ModuleName -> String -> m GHC.DataCon
+lookupDataCon moduleName dataConName = do
+  mapGhcExceptions NotAllowed $ do
+    runGhc $ do
+      GHC.lookupDataCon moduleName dataConName >>= \case
+        Right dataCon -> do
+          pure dataCon
+        Left err -> do
+          throwM $ UnknownError err
+
 moduleIsLoaded :: MonadInterpreter m => ModuleName -> m Bool
 moduleIsLoaded mn = (True <$ findModule mn)
                    `catchIE` (\e -> case e of
diff --git a/src/Hint/GHC.hs b/src/Hint/GHC.hs
index 9b73740..f4a2568 100644
--- a/src/Hint/GHC.hs
+++ b/src/Hint/GHC.hs
@@ -31,6 +31,8 @@ module Hint.GHC (
 #if MIN_VERSION_ghc(9,6,0)
     getPrintUnqual,
 #endif
+    lookupTyCon,
+    lookupDataCon,
     -- * Re-exports
     module X,
 ) where
@@ -135,9 +137,15 @@ import ConLike as X (ConLike(RealDataCon))
 {-------------------- Imports for Shims --------------------}
 
 import Control.Monad.IO.Class (MonadIO)
+
 -- guessTarget
 import qualified GHC (guessTarget)
 
+-- lookupTyCon, lookupDataCon
+import Control.Monad.Trans.Class (lift)
+import Control.Monad.Trans.Writer.CPS (execWriterT, tell)
+import Data.Foldable (for_)
+
 #if MIN_VERSION_ghc(9,6,0)
 -- dynamicGhc
 import GHC.Platform.Ways (hostIsDynamic)
@@ -653,3 +661,45 @@ guessTarget = GHC.guessTarget
 getPrintUnqual :: GhcMonad m => m PrintUnqualified
 getPrintUnqual = GHC.getNamePprCtx
 #endif
+
+lookupCandidates :: GhcMonad m => String -> String -> m [Name]
+lookupCandidates moduleString name = do
+  parseName (moduleString ++ "." ++ name)
+
+lookupTyCon :: GhcMonad m => String -> String -> m (Either String TyCon)
+lookupTyCon moduleString tyConString = do
+  names <- lookupCandidates moduleString tyConString
+  tyCons <- execWriterT $ do
+    for_ names $ \name -> do
+      maybeTyThing <- lift $ lookupName name
+      case maybeTyThing of
+        Just (ATyCon tyCon) -> do
+          tell [tyCon]
+        _ -> do
+          pure ()
+  case tyCons of
+    [tyCon] -> do
+      pure $ Right tyCon
+    [] -> do
+      pure $ Left $ "Could not find type constructor " ++ tyConString
+    _ -> do
+      pure $ Left $ "Ambiguous type constructor " ++ tyConString
+
+lookupDataCon :: GhcMonad m => String -> String -> m (Either String DataCon)
+lookupDataCon moduleString dataConString = do
+  names <- lookupCandidates moduleString dataConString
+  dataCons <- execWriterT $ do
+    for_ names $ \name -> do
+      maybeTyThing <- lift $ lookupName name
+      case maybeTyThing of
+        Just (AConLike (RealDataCon dataCon)) -> do
+          tell [dataCon]
+        _ -> do
+          pure ()
+  case dataCons of
+    [dataCon] -> do
+      pure $ Right dataCon
+    [] -> do
+      pure $ Left $ "Could not find data constructor " ++ dataConString
+    _ -> do
+      pure $ Left $ "Ambiguous data constructor " ++ dataConString
\ No newline at end of file