{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE LambdaCase #-}

-- | The Ollama plugin for GHC
module GHC.Plugin.OllamaHoles where

import Control.Monad (filterM, unless, when)
import Data.Char (isSpace)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.IO qualified as T
import GHC.Plugins hiding ((<>))
import GHC.Tc.Types
import GHC.Tc.Types.Constraint (Hole (..))
import GHC.Tc.Utils.Monad (getGblEnv, newTcRef, writeTcRef, readTcRef)

import GHC.Plugin.OllamaHoles.Backend
import GHC.Plugin.OllamaHoles.Backend.Gemini (geminiBackend)
import GHC.Plugin.OllamaHoles.Backend.Ollama (ollamaBackend)
import GHC.Plugin.OllamaHoles.Backend.OpenAI (openAICompatibleBackend)

import Control.Monad.Catch (handleAll)
import GHC (GhcPs, LHsExpr, GhcRn)
import GHC.Data.StringBuffer qualified as GHC (stringToStringBuffer)
import GHC.Driver.Config.Parser qualified as GHC (initParserOpts)
import GHC.Parser qualified as GHC (parseExpression)
import GHC.Parser.Lexer qualified as GHC (ParseResult (..), getPsErrorMessages, initParserState, unP)
import GHC.Parser.PostProcess qualified as GHC (runPV, unECP)
import GHC.Rename.Expr qualified as GHC (rnLExpr)
import GHC.Tc.Errors.Hole qualified as GHC (tcCheckHoleFit, withoutUnification)
import GHC.Tc.Gen.App qualified as GHC (tcInferSigma)
import GHC.Tc.Utils.TcType qualified as GHC (tyCoFVsOfType)
import GHC.Types.SrcLoc qualified as GHC (mkRealSrcLoc)
#if __GLASGOW_HASKELL__ >= 908
import GHC.Tc.Types.Constraint (CtLocEnv(..))
#endif

#if __GLASGOW_HASKELL__ >= 912
import GHC.Tc.Types.CtLoc (ctLocSpan)
import qualified Data.Map as Map
#else
import GHC.Tc.Types.Constraint (ctLocSpan)
#endif

import Data.Maybe (mapMaybe)

import Data.List (find)
import GHC.Core.TyCo.Rep
import qualified GHC.HsToCore.Docs as GHC
import qualified GHC.Types.Unique.Map as GHC
import qualified GHC.Hs.Doc as GHC
import qualified GHC.Iface.Load as GHC

-- | Prompt used to prompt the LLM
promptTemplate :: Text
promptTemplate :: Text
promptTemplate =
           Text
"Preliminaries:"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"{docs}\n\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"--------------------------------------------------------------------\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"You are a typed-hole plugin within GHC, the Glasgow Haskell Compiler.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"You are given a hole in a Haskell program, and you need to fill it in.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"The hole is represented by the following information:\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"{module}\n{location}\n{imports}\n{hole_var}\n{hole_type}\n{relevant_constraints}\n{candidate_fits}\n\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"{scope}\n\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"{guidance}\n\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Provide one or more Haskell expressions that could fill this hole.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"This means coming up with an expression of the correct type that satisfies the constraints.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Pay special attention to the type of the hole, specifically whether it is a function.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Make sure you synthesize an expression that matches the type of the hole.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Output ONLY the raw Haskell expression(s), one per line.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Do not include explanations, introductions, or any surrounding text.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"If you are using a function from scope, make sure to use the fully qualified name.\n"
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Output a maximum of {numexpr} expresssions.\n"

-- | Determine which backend to use
getBackend :: Flags -> Backend
getBackend :: Flags -> Backend
getBackend Flags{backend_name :: Flags -> Text
backend_name = Text
"ollama"} = Backend
ollamaBackend
getBackend Flags{backend_name :: Flags -> Text
backend_name = Text
"gemini"} = Backend
geminiBackend
getBackend Flags{backend_name :: Flags -> Text
backend_name = Text
"openai", Bool
Int
Text
model_name :: Text
num_expr :: Int
debug :: Bool
include_docs :: Bool
openai_base_url :: Text
openai_key_name :: Text
model_name :: Flags -> Text
num_expr :: Flags -> Int
debug :: Flags -> Bool
include_docs :: Flags -> Bool
openai_base_url :: Flags -> Text
openai_key_name :: Flags -> Text
..} = Text -> Text -> Backend
openAICompatibleBackend Text
openai_base_url Text
openai_key_name
getBackend Flags{Bool
Int
Text
backend_name :: Flags -> Text
model_name :: Flags -> Text
num_expr :: Flags -> Int
debug :: Flags -> Bool
include_docs :: Flags -> Bool
openai_base_url :: Flags -> Text
openai_key_name :: Flags -> Text
model_name :: Text
backend_name :: Text
num_expr :: Int
debug :: Bool
include_docs :: Bool
openai_base_url :: Text
openai_key_name :: Text
..} = String -> Backend
forall a. HasCallStack => String -> a
error (String -> Backend) -> String -> Backend
forall a b. (a -> b) -> a -> b
$ String
"unknown backend: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
backend_name

-- | Ollama plugin for GHC
plugin :: Plugin
plugin :: Plugin
plugin =
    Plugin
defaultPlugin
        { holeFitPlugin = \[String]
opts ->
            HoleFitPluginR -> Maybe HoleFitPluginR
forall a. a -> Maybe a
Just (HoleFitPluginR -> Maybe HoleFitPluginR)
-> HoleFitPluginR -> Maybe HoleFitPluginR
forall a b. (a -> b) -> a -> b
$
                HoleFitPluginR
                    { hfPluginInit :: TcM (TcRef [HoleFitCandidate])
hfPluginInit = [HoleFitCandidate] -> TcM (TcRef [HoleFitCandidate])
forall a gbl lcl. a -> TcRnIf gbl lcl (TcRef a)
newTcRef []
                    , hfPluginStop :: TcRef [HoleFitCandidate] -> TcM ()
hfPluginStop = \TcRef [HoleFitCandidate]
_ -> () -> TcM ()
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                    , hfPluginRun :: TcRef [HoleFitCandidate] -> HoleFitPlugin
hfPluginRun = \TcRef [HoleFitCandidate]
ref ->
                            HoleFitPlugin
                                { candPlugin :: CandPlugin
candPlugin = \TypedHole
_ [HoleFitCandidate]
c -> TcRef [HoleFitCandidate] -> [HoleFitCandidate] -> TcM ()
forall a gbl lcl. TcRef a -> a -> TcRnIf gbl lcl ()
writeTcRef TcRef [HoleFitCandidate]
ref [HoleFitCandidate]
c TcM ()
-> IOEnv (Env TcGblEnv TcLclEnv) [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) [HoleFitCandidate]
forall a b.
IOEnv (Env TcGblEnv TcLclEnv) a
-> IOEnv (Env TcGblEnv TcLclEnv) b
-> IOEnv (Env TcGblEnv TcLclEnv) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) [HoleFitCandidate]
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [HoleFitCandidate]
c
                                , fitPlugin :: FitPlugin
fitPlugin = [String] -> TcRef [HoleFitCandidate] -> FitPlugin
fitPlugin [String]
opts TcRef [HoleFitCandidate]
ref
                                }
                    }
        }
  where
        
    pluginName :: Text
pluginName = Text
"Ollama Plugin"
       
    fitPlugin :: [String] -> TcRef [HoleFitCandidate] -> FitPlugin
fitPlugin [String]
opts TcRef [HoleFitCandidate]
ref TypedHole
hole [HoleFit]
fits = do
        [HoleFitCandidate]
cands <- TcRef [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) [HoleFitCandidate]
forall a gbl lcl. TcRef a -> TcRnIf gbl lcl a
readTcRef TcRef [HoleFitCandidate]
ref
        let flags :: Flags
flags@Flags{Bool
Int
Text
backend_name :: Flags -> Text
model_name :: Flags -> Text
num_expr :: Flags -> Int
debug :: Flags -> Bool
include_docs :: Flags -> Bool
openai_base_url :: Flags -> Text
openai_key_name :: Flags -> Text
model_name :: Text
backend_name :: Text
num_expr :: Int
debug :: Bool
include_docs :: Bool
openai_base_url :: Text
openai_key_name :: Text
..} = [String] -> Flags
parseFlags [String]
opts
        DynFlags
dflags <- IOEnv (Env TcGblEnv TcLclEnv) DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
        TcGblEnv
gbl_env <- TcRnIf TcGblEnv TcLclEnv TcGblEnv
forall gbl lcl. TcRnIf gbl lcl gbl
getGblEnv
        let mod_name :: String
mod_name = ModuleName -> String
moduleNameString (ModuleName -> String) -> ModuleName -> String
forall a b. (a -> b) -> a -> b
$ GenModule Unit -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName (GenModule Unit -> ModuleName) -> GenModule Unit -> ModuleName
forall a b. (a -> b) -> a -> b
$ TcGblEnv -> GenModule Unit
tcg_mod TcGblEnv
gbl_env
            imports :: ImportAvails
imports = TcGblEnv -> ImportAvails
tcg_imports TcGblEnv
gbl_env
        let backend :: Backend
backend = Flags -> Backend
getBackend Flags
flags
        Maybe [Text]
available_models <- IO (Maybe [Text]) -> IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Text])
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe [Text]) -> IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Text]))
-> IO (Maybe [Text])
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe [Text])
forall a b. (a -> b) -> a -> b
$ Backend -> IO (Maybe [Text])
listModels Backend
backend
        case Maybe [Text]
available_models of
            Maybe [Text]
Nothing ->
                String -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit]
forall a. HasCallStack => String -> a
error (String -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit])
-> String -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit]
forall a b. (a -> b) -> a -> b
$
                    String
"--- " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
pluginName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": No models available, check your configuration ---"
            Just [Text]
models -> do
                Bool -> TcM () -> TcM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Text
model_name Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
models) (TcM () -> TcM ()) -> TcM () -> TcM ()
forall a b. (a -> b) -> a -> b
$
                    String -> TcM ()
forall a. HasCallStack => String -> a
error (String -> TcM ()) -> String -> TcM ()
forall a b. (a -> b) -> a -> b
$
                        String
"--- "
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
pluginName
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": Model "
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
model_name
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" not found. "
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ( if Text
backend_name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"ollama"
                                    then String
"Use `ollama pull` to download the model, or "
                                    else String
""
                               )
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"specify another model using "
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"`-fplugin-opt=GHC.Plugin.OllamaHoles:model=<model_name>` ---"
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"--- Availble models: "
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack ([Text] -> Text
T.unlines [Text]
models)
                            String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" ---"
                IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$ Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"--- " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pluginName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": Hole Found ---"
                let mn :: String
mn = String
"Module: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
mod_name
                let lc :: String
lc = String
"Location: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (Maybe RealSrcSpan -> SDoc
forall a. Outputable a => a -> SDoc
ppr (Maybe RealSrcSpan -> SDoc) -> Maybe RealSrcSpan -> SDoc
forall a b. (a -> b) -> a -> b
$ CtLoc -> RealSrcSpan
ctLocSpan (CtLoc -> RealSrcSpan) -> (Hole -> CtLoc) -> Hole -> RealSrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Hole -> CtLoc
hole_loc (Hole -> RealSrcSpan) -> Maybe Hole -> Maybe RealSrcSpan
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypedHole -> Maybe Hole
th_hole TypedHole
hole)
#if __GLASGOW_HASKELL__ >= 912
                let im = "Imports: " <> showSDoc dflags (ppr $ Map.keys $ imp_mods imports)
#else
                let im :: String
im = String
"Imports: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags ([GenModule Unit] -> SDoc
forall a. Outputable a => a -> SDoc
ppr ([GenModule Unit] -> SDoc) -> [GenModule Unit] -> SDoc
forall a b. (a -> b) -> a -> b
$ ModuleEnv [ImportedBy] -> [GenModule Unit]
forall a. ModuleEnv a -> [GenModule Unit]
moduleEnvKeys (ModuleEnv [ImportedBy] -> [GenModule Unit])
-> ModuleEnv [ImportedBy] -> [GenModule Unit]
forall a b. (a -> b) -> a -> b
$ ImportAvails -> ModuleEnv [ImportedBy]
imp_mods ImportAvails
imports)
#endif
                case TypedHole -> Maybe Hole
th_hole TypedHole
hole of
                    Just Hole
h -> do
                        let hv :: String
hv = String
"Hole variable: _" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> OccName -> String
occNameString (RdrName -> OccName
forall name. HasOccName name => name -> OccName
occName (RdrName -> OccName) -> RdrName -> OccName
forall a b. (a -> b) -> a -> b
$ Hole -> RdrName
hole_occ Hole
h)
                        let ht :: String
ht = String
"Hole type: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (KindOrType -> SDoc
forall a. Outputable a => a -> SDoc
ppr (KindOrType -> SDoc) -> KindOrType -> SDoc
forall a b. (a -> b) -> a -> b
$ Hole -> KindOrType
hole_ty Hole
h)
                        let rc :: String
rc = String
"Relevant constraints: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (Bag CtEvidence -> SDoc
forall a. Outputable a => a -> SDoc
ppr (Bag CtEvidence -> SDoc) -> Bag CtEvidence -> SDoc
forall a b. (a -> b) -> a -> b
$ TypedHole -> Bag CtEvidence
th_relevant_cts TypedHole
hole)
                        let cf :: String
cf = String
"Candidate fits: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags ([HoleFit] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [HoleFit]
fits)
                        let scope :: String
scope = String
"Things in scope: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags ([RdrName] -> SDoc
forall a. Outputable a => a -> SDoc
ppr ([RdrName] -> SDoc) -> [RdrName] -> SDoc
forall a b. (a -> b) -> a -> b
$ (HoleFitCandidate -> Maybe RdrName)
-> [HoleFitCandidate] -> [RdrName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe HoleFitCandidate -> Maybe RdrName
fullyQualified [HoleFitCandidate]
cands)
                        String
docs <- if Bool
include_docs then [HoleFitCandidate] -> TcM String
getDocs [HoleFitCandidate]
cands else String -> TcM String
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return String
""
  
                        String
guide <- [HoleFitCandidate] -> TcM String
seekGuidance [HoleFitCandidate]
cands
                        let prompt' :: Text
prompt' =
                                Text -> [(Text, String)] -> Text
replacePlaceholders
                                    Text
promptTemplate
                                    [ (Text
"{module}", String
mn)
                                    , (Text
"{location}", String
lc)
                                    , (Text
"{imports}", String
im)
                                    , (Text
"{hole_var}", String
hv)
                                    , (Text
"{hole_type}", String
ht)
                                    , (Text
"{relevant_constraints}", String
rc)
                                    , (Text
"{candidate_fits}", String
cf)
                                    , (Text
"{numexpr}", Int -> String
forall a. Show a => a -> String
show Int
num_expr)
                                    , (Text
"{guidance}", String
guide)
                                    , (Text
"{scope}", String
scope)
                                    , (Text
"{docs}", String
docs)
                                    ]
                        Either String Text
res <- IO (Either String Text)
-> IOEnv (Env TcGblEnv TcLclEnv) (Either String Text)
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either String Text)
 -> IOEnv (Env TcGblEnv TcLclEnv) (Either String Text))
-> IO (Either String Text)
-> IOEnv (Env TcGblEnv TcLclEnv) (Either String Text)
forall a b. (a -> b) -> a -> b
$ Backend -> Text -> Text -> IO (Either String Text)
generateFits Backend
backend Text
prompt' Text
model_name
                        case Either String Text
res of
                            Right Text
rsp -> do
                                let lns :: [Text]
lns = ([Text] -> [Text]
preProcess ([Text] -> [Text]) -> (Text -> [Text]) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines) Text
rsp
                                [Text]
verified <- (Text -> IOEnv (Env TcGblEnv TcLclEnv) Bool)
-> [Text] -> IOEnv (Env TcGblEnv TcLclEnv) [Text]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (Bool -> TypedHole -> Text -> IOEnv (Env TcGblEnv TcLclEnv) Bool
verifyHoleFit Bool
debug TypedHole
hole) [Text]
lns
                                IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$ Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                                    Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"--- " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pluginName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": Prompt ---\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
prompt'
                                    Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"--- " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
pluginName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": Response ---\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
rsp
                                let fits' :: [HoleFit]
fits' = (Text -> HoleFit) -> [Text] -> [HoleFit]
forall a b. (a -> b) -> [a] -> [b]
map (SDoc -> HoleFit
RawHoleFit (SDoc -> HoleFit) -> (Text -> SDoc) -> Text -> HoleFit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> SDoc
forall doc. IsLine doc => String -> doc
text (String -> SDoc) -> (Text -> String) -> Text -> SDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) [Text]
verified
                                -- Return the generated fits
                                [HoleFit] -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit]
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [HoleFit]
fits'
                            Left String
err -> do
                                IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$
                                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                                        Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$
                                            Text
pluginName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" failed to generate a response.\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
err
                                -- Return the original fits without modification
                                [HoleFit] -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit]
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [HoleFit]
fits
                    Maybe Hole
Nothing -> [HoleFit] -> IOEnv (Env TcGblEnv TcLclEnv) [HoleFit]
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return [HoleFit]
fits


-- | Parse an expression in the current context
parseInContext :: Text -> TcM (Either String (LHsExpr GhcPs))
parseInContext :: Text -> TcM (Either String (LHsExpr GhcPs))
parseInContext Text
fit = do
    DynFlags
dflags <- IOEnv (Env TcGblEnv TcLclEnv) DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    let parsed :: ParseResult (LocatedA (HsExpr GhcPs))
parsed =
            P (LocatedA (HsExpr GhcPs))
-> PState -> ParseResult (LocatedA (HsExpr GhcPs))
forall a. P a -> PState -> ParseResult a
GHC.unP (P ECP
GHC.parseExpression P ECP
-> (ECP -> P (LocatedA (HsExpr GhcPs)))
-> P (LocatedA (HsExpr GhcPs))
forall a b. P a -> (a -> P b) -> P b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ECP
p -> PV (LocatedA (HsExpr GhcPs)) -> P (LocatedA (HsExpr GhcPs))
forall a. PV a -> P a
GHC.runPV (PV (LocatedA (HsExpr GhcPs)) -> P (LocatedA (HsExpr GhcPs)))
-> PV (LocatedA (HsExpr GhcPs)) -> P (LocatedA (HsExpr GhcPs))
forall a b. (a -> b) -> a -> b
$ ECP -> forall b. DisambECP b => PV (LocatedA b)
GHC.unECP ECP
p) (PState -> ParseResult (LocatedA (HsExpr GhcPs)))
-> PState -> ParseResult (LocatedA (HsExpr GhcPs))
forall a b. (a -> b) -> a -> b
$
                ParserOpts -> StringBuffer -> RealSrcLoc -> PState
GHC.initParserState
                    (DynFlags -> ParserOpts
GHC.initParserOpts DynFlags
dflags)
                    (String -> StringBuffer
GHC.stringToStringBuffer (Text -> String
T.unpack Text
fit))
                    (FastString -> Int -> Int -> RealSrcLoc
GHC.mkRealSrcLoc (String -> FastString
mkFastString String
"<hole-fit-validation>") Int
1 Int
1)
    case ParseResult (LocatedA (HsExpr GhcPs))
parsed of
      GHC.PFailed PState
st -> Either String (LocatedA (HsExpr GhcPs))
-> IOEnv
     (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs)))
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (LocatedA (HsExpr GhcPs))
 -> IOEnv
      (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs))))
-> Either String (LocatedA (HsExpr GhcPs))
-> IOEnv
     (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs)))
forall a b. (a -> b) -> a -> b
$ String -> Either String (LocatedA (HsExpr GhcPs))
forall a b. a -> Either a b
Left (Messages PsMessage -> String
forall a. Outputable a => a -> String
showPprUnsafe (Messages PsMessage -> String) -> Messages PsMessage -> String
forall a b. (a -> b) -> a -> b
$ PState -> Messages PsMessage
GHC.getPsErrorMessages PState
st)
      GHC.POk PState
_ (LHsExpr GhcPs
p_e :: LHsExpr GhcPs) -> Either String (LocatedA (HsExpr GhcPs))
-> IOEnv
     (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs)))
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (LocatedA (HsExpr GhcPs))
 -> IOEnv
      (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs))))
-> Either String (LocatedA (HsExpr GhcPs))
-> IOEnv
     (Env TcGblEnv TcLclEnv) (Either String (LocatedA (HsExpr GhcPs)))
forall a b. (a -> b) -> a -> b
$ LocatedA (HsExpr GhcPs) -> Either String (LocatedA (HsExpr GhcPs))
forall a b. b -> Either a b
Right LHsExpr GhcPs
LocatedA (HsExpr GhcPs)
p_e
      

-- | Check that the hole fit matches the type of the hole
verifyHoleFit :: Bool -> TypedHole -> Text -> TcM Bool
verifyHoleFit :: Bool -> TypedHole -> Text -> IOEnv (Env TcGblEnv TcLclEnv) Bool
verifyHoleFit Bool
debug TypedHole
hole Text
fit | Just Hole
h <- TypedHole -> Maybe Hole
th_hole TypedHole
hole = (SomeException -> IOEnv (Env TcGblEnv TcLclEnv) Bool)
-> IOEnv (Env TcGblEnv TcLclEnv) Bool
-> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall (m :: * -> *) a.
(HasCallStack, MonadCatch m) =>
(SomeException -> m a) -> m a -> m a
handleAll SomeException -> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall {m :: * -> *} {a}. (MonadIO m, Show a) => a -> m Bool
falseOnErr (IOEnv (Env TcGblEnv TcLclEnv) Bool
 -> IOEnv (Env TcGblEnv TcLclEnv) Bool)
-> IOEnv (Env TcGblEnv TcLclEnv) Bool
-> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall a b. (a -> b) -> a -> b
$ do
    -- Instaniate a new IORef session with the current HscEnv.
    -- First we try parsing the suggest hole-fits
    Either String (LocatedA (HsExpr GhcPs))
parsed <- Text -> TcM (Either String (LHsExpr GhcPs))
parseInContext Text
fit
    case Either String (LocatedA (HsExpr GhcPs))
parsed of
        Left String
err_msg-> do
            IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$ Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              String -> IO ()
putStrLn String
"--- Error when validating: ---"
              String -> IO ()
putStrLn String
err_msg
            Bool -> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        Right LocatedA (HsExpr GhcPs)
p_e -> do
            -- If parsing was successful, we try renaming the expression
            (GenLocated SrcSpanAnnA (HsExpr GhcRn)
rn_e, FreeVars
free_vars) <- LHsExpr GhcPs -> RnM (LHsExpr GhcRn, FreeVars)
GHC.rnLExpr LHsExpr GhcPs
LocatedA (HsExpr GhcPs)
p_e
            Bool -> TcM () -> TcM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (TcM () -> TcM ()) -> TcM () -> TcM ()
forall a b. (a -> b) -> a -> b
$ IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ FreeVars -> String
forall a. Outputable a => a -> String
showPprUnsafe FreeVars
free_vars
            -- Finally, we infer the type of the expression
            KindOrType
expr_ty <- Bool -> LHsExpr GhcRn -> TcM KindOrType
GHC.tcInferSigma Bool
False LHsExpr GhcRn
GenLocated SrcSpanAnnA (HsExpr GhcRn)
rn_e
            -- And check whether that is indeed a valid hole fit
            (Bool
does_fit, HsWrapper
wrapper) <-
                FV -> TcM (Bool, HsWrapper) -> TcM (Bool, HsWrapper)
forall a. FV -> TcM a -> TcM a
GHC.withoutUnification (KindOrType -> FV
GHC.tyCoFVsOfType (KindOrType -> FV) -> KindOrType -> FV
forall a b. (a -> b) -> a -> b
$ Hole -> KindOrType
hole_ty Hole
h) (TcM (Bool, HsWrapper) -> TcM (Bool, HsWrapper))
-> TcM (Bool, HsWrapper) -> TcM (Bool, HsWrapper)
forall a b. (a -> b) -> a -> b
$
                    TypedHole -> KindOrType -> KindOrType -> TcM (Bool, HsWrapper)
GHC.tcCheckHoleFit TypedHole
hole (Hole -> KindOrType
hole_ty Hole
h) KindOrType
expr_ty
            Bool -> TcM () -> TcM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (TcM () -> TcM ()) -> TcM () -> TcM ()
forall a b. (a -> b) -> a -> b
$ IO () -> TcM ()
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TcM ()) -> IO () -> TcM ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ HsWrapper -> String
forall a. Outputable a => a -> String
showPprUnsafe HsWrapper
wrapper
            Bool -> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
does_fit
  where
    falseOnErr :: a -> m Bool
falseOnErr a
e = IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ do
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            String -> IO ()
putStrLn String
"--- Error when validating: ---"
            a -> IO ()
forall a. Show a => a -> IO ()
print a
e
        Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
verifyHoleFit Bool
_ TypedHole
_ Text
_ = Bool -> IOEnv (Env TcGblEnv TcLclEnv) Bool
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False


-- | Try to find the guide provided by the user
seekGuidance :: [HoleFitCandidate] -> TcM String
seekGuidance :: [HoleFitCandidate] -> TcM String
seekGuidance [HoleFitCandidate]
cands = do
    case (HoleFitCandidate -> Bool)
-> [HoleFitCandidate] -> Maybe HoleFitCandidate
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"_guide") (String -> Bool)
-> (HoleFitCandidate -> String) -> HoleFitCandidate -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
forall a. Outputable a => a -> String
showPprUnsafe (OccName -> String)
-> (HoleFitCandidate -> OccName) -> HoleFitCandidate -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HoleFitCandidate -> OccName
forall name. HasOccName name => name -> OccName
occName) [HoleFitCandidate]
cands of
      Just (IdHFCand Id
i) | KindOrType
ty <- Id -> KindOrType
idType Id
i -> do
          case KindOrType
ty of
            TyConApp TyCon
tc [KindOrType
errm, KindOrType
errm_t] | String
"Proxy" <- TyCon -> String
forall a. Outputable a => a -> String
showPprUnsafe TyCon
tc,
                                         String
"ErrorMessage" <- KindOrType -> String
forall a. Outputable a => a -> String
showPprUnsafe KindOrType
errm,
                                         TyConApp TyCon
_ [KindOrType
guide_msg] <- KindOrType
errm_t ->
              String -> TcM String
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> TcM String) -> String -> TcM String
forall a b. (a -> b) -> a -> b
$ String
"The user provided these instructions: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> KindOrType -> String
forall a. Outputable a => a -> String
showPprUnsafe KindOrType
guide_msg
            KindOrType
_ -> String -> TcM String
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return String
""
      Maybe HoleFitCandidate
_ -> String -> TcM String
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return String
""

-- | Preprocess the response to remove empty lines, lines with only spaces, and code blocks
preProcess :: [Text] -> [Text]
preProcess :: [Text] -> [Text]
preProcess [] = []
-- \| Remove lines between <think> and </think> tags from e.g. deepseek
preProcess (Text
ln : [Text]
lns)
    | Text -> Text -> Bool
T.isPrefixOf Text
"<think>" Text
ln =
        [Text] -> [Text]
preProcess (Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
drop Int
1 ([Text] -> [Text]) -> [Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (Text -> Bool) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text -> Bool
T.isPrefixOf Text
"</think>") [Text]
lns)
preProcess (Text
ln : [Text]
lns) | Bool
should_drop = [Text] -> [Text]
preProcess [Text]
lns
  where
    should_drop :: Bool
    should_drop :: Bool
should_drop =
        Text -> Bool
T.null Text
ln
            Bool -> Bool -> Bool
|| (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isSpace Text
ln
            Bool -> Bool -> Bool
|| Text -> Text -> Bool
T.isPrefixOf Text
"```" Text
ln
preProcess (Text
ln : [Text]
lns) = Text -> Text
transform Text
ln Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [Text] -> [Text]
preProcess [Text]
lns
  where
    transform :: Text -> Text
    transform :: Text -> Text
transform = Text -> Text
T.strip

-- | Command line options for the plugin
data Flags = Flags
    { Flags -> Text
model_name :: Text
    , Flags -> Text
backend_name :: Text
    , Flags -> Int
num_expr :: Int
    , Flags -> Bool
debug :: Bool
    , Flags -> Bool
include_docs :: Bool
    , Flags -> Text
openai_base_url :: Text
    , Flags -> Text
openai_key_name :: Text
    }

-- | Default flags for the plugin
defaultFlags :: Flags
defaultFlags :: Flags
defaultFlags =
    Flags
        { model_name :: Text
model_name = Text
"gemma3:27b-it-qat"
        , backend_name :: Text
backend_name = Text
"ollama"
        , num_expr :: Int
num_expr = Int
5
        , debug :: Bool
debug = Bool
False
        , include_docs :: Bool
include_docs = Bool
False
        , openai_base_url :: Text
openai_base_url = Text
"https://api.openai.com"
        , openai_key_name :: Text
openai_key_name = Text
"OPENAI_API_KEY"
        }

-- | Produce the documentation of all the HolefitCandidates.
getDocs :: [HoleFitCandidate] -> TcM String
getDocs :: [HoleFitCandidate] -> TcM String
getDocs [HoleFitCandidate]
cs = do
    DynFlags
dflags <- IOEnv (Env TcGblEnv TcLclEnv) DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
    TcGblEnv
gbl_env <- TcRnIf TcGblEnv TcLclEnv TcGblEnv
forall gbl lcl. TcRnIf gbl lcl gbl
getGblEnv
    UniqMap Name [HsDoc GhcRn]
lcl_docs <-  (Maybe Docs -> UniqMap Name [HsDoc GhcRn])
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe Docs)
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
forall a b.
(a -> b)
-> IOEnv (Env TcGblEnv TcLclEnv) a
-> IOEnv (Env TcGblEnv TcLclEnv) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ( UniqMap Name [HsDoc GhcRn]
-> (Docs -> UniqMap Name [HsDoc GhcRn])
-> Maybe Docs
-> UniqMap Name [HsDoc GhcRn]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe UniqMap Name [HsDoc GhcRn]
forall k a. UniqMap k a
GHC.emptyUniqMap Docs -> UniqMap Name [HsDoc GhcRn]
GHC.docs_decls) (IOEnv (Env TcGblEnv TcLclEnv) (Maybe Docs)
 -> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn]))
-> (IO (Maybe Docs) -> IOEnv (Env TcGblEnv TcLclEnv) (Maybe Docs))
-> IO (Maybe Docs)
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Maybe Docs) -> IOEnv (Env TcGblEnv TcLclEnv) (Maybe Docs)
forall a. IO a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Docs)
 -> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn]))
-> IO (Maybe Docs)
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
forall a b. (a -> b) -> a -> b
$ DynFlags -> TcGblEnv -> IO (Maybe Docs)
forall (m :: * -> *).
MonadIO m =>
DynFlags -> TcGblEnv -> m (Maybe Docs)
GHC.extractDocs DynFlags
dflags TcGblEnv
gbl_env
    UniqMap Name [HsDoc GhcRn]
all_docs <- UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs UniqMap Name [HsDoc GhcRn]
lcl_docs [HoleFitCandidate]
cs
    String -> TcM String
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> TcM String) -> String -> TcM String
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (String -> Bool) -> String -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ (HoleFitCandidate -> String) -> [HoleFitCandidate] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (DynFlags
-> UniqMap Name [HsDoc GhcRn] -> HoleFitCandidate -> String
mkDoc DynFlags
dflags UniqMap Name [HsDoc GhcRn]
all_docs) [HoleFitCandidate]
cs
 where
    allDocs :: UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs UniqMap Name [HsDoc GhcRn]
docs (IdHFCand Id
_:[HoleFitCandidate]
cs') = UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs UniqMap Name [HsDoc GhcRn]
docs [HoleFitCandidate]
cs'
    allDocs UniqMap Name [HsDoc GhcRn]
docs (GreHFCand GlobalRdrElt
gre:[HoleFitCandidate]
cs') | GlobalRdrElt -> Bool
gre_lcl GlobalRdrElt
gre = UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs UniqMap Name [HsDoc GhcRn]
docs [HoleFitCandidate]
cs'
    allDocs UniqMap Name [HsDoc GhcRn]
docs (HoleFitCandidate
c:[HoleFitCandidate]
cs') = do
        Maybe Docs
if_docs <- ModIface_ 'ModIfaceFinal -> Maybe Docs
forall (phase :: ModIfacePhase). ModIface_ phase -> Maybe Docs
mi_docs (ModIface_ 'ModIfaceFinal -> Maybe Docs)
-> IOEnv (Env TcGblEnv TcLclEnv) (ModIface_ 'ModIfaceFinal)
-> IOEnv (Env TcGblEnv TcLclEnv) (Maybe Docs)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SDoc
-> Name -> IOEnv (Env TcGblEnv TcLclEnv) (ModIface_ 'ModIfaceFinal)
GHC.loadInterfaceForName (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"hole-fit docs") (HoleFitCandidate -> Name
forall a. NamedThing a => a -> Name
getName HoleFitCandidate
c)
        case Maybe Docs
if_docs of
          Maybe Docs
Nothing -> UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs UniqMap Name [HsDoc GhcRn]
docs [HoleFitCandidate]
cs'
          Just Docs
d -> UniqMap Name [HsDoc GhcRn]
-> [HoleFitCandidate]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
allDocs (UniqMap Name [HsDoc GhcRn]
-> UniqMap Name [HsDoc GhcRn] -> UniqMap Name [HsDoc GhcRn]
forall k a. UniqMap k a -> UniqMap k a -> UniqMap k a
GHC.plusUniqMap UniqMap Name [HsDoc GhcRn]
docs (Docs -> UniqMap Name [HsDoc GhcRn]
GHC.docs_decls Docs
d)) [HoleFitCandidate]
cs'
    allDocs UniqMap Name [HsDoc GhcRn]
docs [] = UniqMap Name [HsDoc GhcRn]
-> IOEnv (Env TcGblEnv TcLclEnv) (UniqMap Name [HsDoc GhcRn])
forall a. a -> IOEnv (Env TcGblEnv TcLclEnv) a
forall (m :: * -> *) a. Monad m => a -> m a
return UniqMap Name [HsDoc GhcRn]
docs
    mkDoc :: DynFlags -> GHC.UniqMap Name [GHC.HsDoc GhcRn] -> HoleFitCandidate -> String
    mkDoc :: DynFlags
-> UniqMap Name [HsDoc GhcRn] -> HoleFitCandidate -> String
mkDoc DynFlags
dflags UniqMap Name [HsDoc GhcRn]
all_docs HoleFitCandidate
hfc | Just RdrName
rn <- HoleFitCandidate -> Maybe RdrName
fullyQualified HoleFitCandidate
hfc,
                                Just [HsDoc GhcRn]
doc <- UniqMap Name [HsDoc GhcRn] -> Name -> Maybe [HsDoc GhcRn]
forall k a. Uniquable k => UniqMap k a -> k -> Maybe a
GHC.lookupUniqMap UniqMap Name [HsDoc GhcRn]
all_docs (HoleFitCandidate -> Name
forall a. NamedThing a => a -> Name
getName HoleFitCandidate
hfc) = 
       String
"Documentation for `" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (RdrName -> SDoc
forall a. Outputable a => a -> SDoc
ppr RdrName
rn ) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"`:\n```\n"
       String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DynFlags -> [HsDoc GhcRn] -> String
processDoc DynFlags
dflags [HsDoc GhcRn]
doc
       String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"```"
    mkDoc DynFlags
_ UniqMap Name [HsDoc GhcRn]
_ HoleFitCandidate
_ = String
""
    -- We get the first paragraph of the docs to avoid too much context
    processDoc :: DynFlags -> [GHC.HsDoc GhcRn] -> String
    processDoc :: DynFlags -> [HsDoc GhcRn] -> String
processDoc DynFlags
dflags [HsDoc GhcRn]
docs = String
first_paragraph
      where whole_string :: String
whole_string = [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (HsDoc GhcRn -> String) -> [HsDoc GhcRn] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (SDoc -> String) -> (HsDoc GhcRn -> SDoc) -> HsDoc GhcRn -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HsDoc GhcRn -> SDoc
forall a. Outputable a => a -> SDoc
ppr) [HsDoc GhcRn]
docs
            first_paragraph :: String
first_paragraph = [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Bool -> Bool
not (Bool -> Bool) -> (String -> Bool) -> String -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ String -> [String]
lines String
whole_string

-- | Produce a fully qualified name, e.g. L.sort if Data.List is imported as L
fullyQualified :: HoleFitCandidate -> Maybe RdrName
fullyQualified :: HoleFitCandidate -> Maybe RdrName
fullyQualified (IdHFCand Id
i) = RdrName -> Maybe RdrName
forall a. a -> Maybe a
Just (RdrName -> Maybe RdrName) -> RdrName -> Maybe RdrName
forall a b. (a -> b) -> a -> b
$ Id -> RdrName
forall thing. NamedThing thing => thing -> RdrName
getRdrName Id
i
fullyQualified (NameHFCand Name
n) = RdrName -> Maybe RdrName
forall a. a -> Maybe a
Just (RdrName -> Maybe RdrName) -> RdrName -> Maybe RdrName
forall a b. (a -> b) -> a -> b
$ Name -> RdrName
forall thing. NamedThing thing => thing -> RdrName
getRdrName Name
n
fullyQualified (GreHFCand GlobalRdrElt
gre) | (RdrName
n:[RdrName]
_) <- GlobalRdrElt -> [RdrName]
greRdrNames GlobalRdrElt
gre = RdrName -> Maybe RdrName
forall a. a -> Maybe a
Just RdrName
n
fullyQualified HoleFitCandidate
_ = Maybe RdrName
forall a. Maybe a
Nothing

-- | Parse command line options
parseFlags :: [CommandLineOption] -> Flags
parseFlags :: [String] -> Flags
parseFlags = Flags -> [String] -> Flags
parseFlags' Flags
defaultFlags
  where
    parseFlags' :: Flags -> [CommandLineOption] -> Flags
    parseFlags' :: Flags -> [String] -> Flags
parseFlags' Flags
flags [] = Flags
flags
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"model=" (String -> Text
T.pack String
opt) =
            let model_name :: Text
model_name = Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"model=") (String -> Text
T.pack String
opt)
             in Flags -> [String] -> Flags
parseFlags' Flags
flags{model_name = model_name} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"backend=" (String -> Text
T.pack String
opt) =
            let backend_name :: Text
backend_name = Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"backend=") (String -> Text
T.pack String
opt)
             in Flags -> [String] -> Flags
parseFlags' Flags
flags{backend_name = backend_name} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"openai_base_url=" (String -> Text
T.pack String
opt) =
            let openai_base_url :: Text
openai_base_url = Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"openai_base_url=") (String -> Text
T.pack String
opt)
             in Flags -> [String] -> Flags
parseFlags' Flags
flags{openai_base_url = openai_base_url} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"openai_key_name=" (String -> Text
T.pack String
opt) =
            let openai_key_name :: Text
openai_key_name = Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"openai_key_name=") (String -> Text
T.pack String
opt)
             in Flags -> [String] -> Flags
parseFlags' Flags
flags{openai_key_name = openai_key_name} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"debug" (String -> Text
T.pack String
opt) = Flags -> [String] -> Flags
parseFlags' Flags
flags{debug = True} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"include-docs" (String -> Text
T.pack String
opt) = Flags -> [String] -> Flags
parseFlags' Flags
flags{include_docs = True} [String]
opts
    parseFlags' Flags
flags (String
opt : [String]
opts)
        | Text -> Text -> Bool
T.isPrefixOf Text
"n=" (String -> Text
T.pack String
opt) =
            let num_expr :: String
num_expr = Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"n=") (String -> Text
T.pack String
opt)
             in Flags -> [String] -> Flags
parseFlags' Flags
flags{num_expr = read num_expr} [String]
opts
    parseFlags' Flags
flags [String]
_ = Flags
flags

-- | Helper function to replace placeholders in a template string
replacePlaceholders :: Text -> [(Text, String)] -> Text
replacePlaceholders :: Text -> [(Text, String)] -> Text
replacePlaceholders = (Text -> (Text, String) -> Text)
-> Text -> [(Text, String)] -> Text
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Text -> (Text, String) -> Text
replacePlaceholder
  where
    replacePlaceholder :: Text -> (Text, String) -> Text
    replacePlaceholder :: Text -> (Text, String) -> Text
replacePlaceholder Text
str (Text
placeholder, String
value) = HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
placeholder (String -> Text
T.pack String
value) Text
str