ecta-plugin-0.1.1.3: ECTA/Plugin.hs
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
module ECTA.Plugin (plugin) where
import GhcPlugins hiding ((<>))
import TcHoleErrors
import TcHoleFitTypes
import TcRnTypes
import Constraint
import ECTA.Plugin.Utils
import Application.TermSearch.Dataset
import Application.TermSearch.Type
import Application.TermSearch.TermSearch hiding (allConstructors, generalize)
import Data.ECTA
import Data.ECTA.Term
import qualified Data.Map as Map
import Data.Text (pack, unpack, Text)
import Data.Maybe (fromMaybe, mapMaybe, isJust, fromJust)
import Data.Tuple (swap)
import Data.List (sortOn, groupBy, nub, nubBy, (\\))
import Data.Function (on)
import qualified Data.Monoid as M
import MonadUtils (concatMapM)
import TcRnMonad (writeTcRef, newTcRef, readTcRef, mapMaybeM, getTopEnv, tryAllM)
import TcEnv (tcLookupId, tcLookupIdMaybe, tcLookup)
import qualified Data.Bifunctor as Bi
import TcRnDriver (tcRnGetInfo)
import GHC (ClsInst)
import InstEnv (ClsInst(ClsInst, is_tvs, is_cls_nm, is_tys), is_dfun)
import ConLike (ConLike(RealDataCon))
import Data.ECTA.Paths (Path, mkEqConstraints)
import Application.TermSearch.Utils
import Data.Containers.ListUtils (nubOrd)
import Debug.Trace
import Data.Either (partitionEithers)
import Text.Read (readMaybe)
import qualified Data.Set as Set
import Control.Exception (evaluate, displayException)
import System.IO (hPutStrLn, stderr)
plugin :: Plugin
plugin =
defaultPlugin
{ holeFitPlugin = \opts ->
Just $
HoleFitPluginR {
hfPluginInit = newTcRef [],
hfPluginRun = \ref ->
( HoleFitPlugin
{ candPlugin = \_ c -> writeTcRef ref c >> return c,
fitPlugin = \h fts -> readTcRef ref >>= ectaPlugin opts h fts
}
),
hfPluginStop = const (return ())
}
}
candsToComps :: [HoleFitCandidate] -> TcM [((Either Text Text, TypeSkeleton), [Type])]
candsToComps = mapMaybeM (fmap (fmap extract) . candToTN)
where candToTN :: HoleFitCandidate -> TcM (Maybe (Either Text Text, (TypeSkeleton, [Type])))
candToTN cand = fmap (fmap (nm,) . (>>= typeToSkeleton)) (c2t cand)
where nm = (case cand of
IdHFCand _ -> Left
_ -> Right) $ pack $ occNameString $ occName cand
c2t cand =
case cand of
IdHFCand id -> return $ Just $ idType id
NameHFCand nm -> tcTyThingTypeMaybe <$> tcLookup nm
GreHFCand GRE{..} -> tcTyThingTypeMaybe <$> tcLookup gre_name
extract (a, (b,c)) = ((a,b), c)
tcTyThingTypeMaybe :: TcTyThing -> Maybe Type
tcTyThingTypeMaybe (ATcId tttid _) = Just $ idType tttid
tcTyThingTypeMaybe (AGlobal (AnId ttid)) =Just $ idType ttid
tcTyThingTypeMaybe (AGlobal (ATyCon ttid)) | t <- mkTyConApp ttid [],
(tcReturnsConstraintKind . tcTypeKind) t
= Just t
tcTyThingTypeMaybe (AGlobal (AConLike (RealDataCon con))) = Just $ idType $ dataConWorkId con
tcTyThingTypeMaybe _ = Nothing
instToTerm :: ClsInst -> Maybe (Text, TypeSkeleton)
instToTerm ClsInst{..} | -- length is_tvs <= 1, -- uncomment if you want explosion!
Just (tyskel,args) <- typeToSkeleton $ idType is_dfun
= Just (toDictStr $ clsstr <> tystr, tyskel )
where clsstr = pack $ showSDocUnsafe $ ppr is_cls_nm
tystr = pack $ showSDocUnsafe $ ppr is_tys
instToTerm _ = Nothing
toDictStr :: Text -> Text
toDictStr t = spToUnderscore $ "<@" <> t <> "@>"
spToUnderscore :: Text -> Text
spToUnderscore = pack . sp . unpack
where sp (' ':str) = '_':sp str
sp (s:str) = s:sp str
sp [] = []
defaultSize :: Int
defaultSize = 5
-- | Parses the options and returns the max expression size to use.
-- limited to 5 by default.
getExprSize :: [CommandLineOption] -> Int
getExprSize (o:opts) | ("expr-size",'=':n) <- span (/= '=') o,
Just x <- readMaybe n = x
getExprSize _ = defaultSize
dedup :: [Text] -> [Text]
dedup ts = dedup' Set.empty ts
where dedup' seen [] = []
dedup' seen (t:ts) | Set.member t seen = dedup' seen ts
dedup' seen (t:ts) = t:dedup' (Set.insert t seen) ts
ectaPlugin :: [CommandLineOption] -> TypedHole
-> [HoleFit] -> [HoleFitCandidate] -> TcM [HoleFit]
ectaPlugin opts TyH{..} found_fits scope | Just hole <- tyHCt,
expr_size <- getExprSize opts,
ty <- ctPred hole = do
let hM act = do io_r <- tryAllM $ do inner_r <- act
-- we have to force the evaluation of
-- the elements in the list, otherwise
-- the error doesn't show up... and then
-- we can't catch it.
_ <- liftIO $ mapM evaluate inner_r
return inner_r
case io_r of
Left e -> do liftIO $ do hPutStrLn stderr $ "Hectare error:"
hPutStrLn stderr $ (displayException e)
return found_fits
Right r -> return r
hM $ do
(fun_comps, scons) <- fmap (nubBy eqType . concat) . unzip <$> candsToComps scope
let (local_comps, global_comps) = partitionEithers $ map to_e fun_comps
to_e (Left t,ts) = Left (t,ts)
to_e (Right t, ts) = Right (t,ts)
-- The constraints are there and added to the graph... but we have to
-- be more precise when we add them to the machine. Any time a
-- function requires a constraint to hold for one of it's variables,
-- we have to add a path equality to the ECTA.
let constraints = filter (tcReturnsConstraintKind . tcTypeKind) scons
hsc_env <- getTopEnv
instance_comps <- mapMaybe instToTerm . concat <$>
mapMaybeM (fmap (fmap (\(_,_,c,_,_) -> c) . snd)
. liftIO . tcRnGetInfo hsc_env . getName
. tyConAppTyCon) constraints
case typeToSkeleton ty of
Just (t, cons) | -- isSafe t,
resNode <- typeToFta t -> do
let givens = concatMap (map idType . ic_given) tyHImplics
g2c g = fmap (toDictStr (pack $ showSDocUnsafe $ ppr g),)
$ fmap fst $ typeToSkeleton g
given_comps = mapMaybe g2c givens
local_scope_comps = local_comps ++ given_comps
global_scope_comps = global_comps ++ instance_comps
scope_comps = local_scope_comps ++ global_scope_comps
-- let (scopeNode, anyArg, argNodes, skels, groups) =
argNodes = ngnodes local_scope_comps
addSyms st tt = map (Bi.bimap (Symbol . st) (tt . typeToFta))
-- . filter (\(_,t) -> isSafe t)
gnodes = addSyms id (generalize global_scope_comps)
ngnodes = addSyms id id
anyArg = Node $ map (\(s,t) -> Edge s [t]) $
(gnodes global_scope_comps) ++ argNodes
scopeNode = anyArg
skels = Map.fromList $ scope_comps
groups = Map.fromList $ map (\(t,_) -> (t,[t])) scope_comps
-- in (scopeNode, anyArg, argNodes, skels, groups)
-- We ignore ppterms for now, because they need to be printed differently.
-- let res = getAllTerms $ refold $ reduceFully $ filterType scopeNode resNode
-- ppterms <- concatMapM (prettyMatch skels groups . prettyTerm ) res
let even_more_terms =
map (ppNoPar . prettyTerm) $
concatMap (getAllTerms . refold . reduceFully . flip filterType resNode )
(rtkUpToKAtLeast1 argNodes scope_comps anyArg True expr_size)
--text_fits = ppterms ++ even_more_terms
ecta_fits = dedup even_more_terms
fit_set = Set.fromList $ mapMaybe f found_fits
where f (HoleFit {hfCand=c}) = Just (pack $ occNameString $ occName c)
f _ = Nothing
filtered_fits = map (RawHoleFit . text . unpack . parIfReq) $
filter (not . flip Set.member fit_set) ecta_fits
return $ found_fits ++ filtered_fits
_ -> return found_fits
-- TODO:
-- 1. I think we need to add type applications, i.e. [] @a or similar, let's see.