diff --git a/ECTA/Plugin.hs b/ECTA/Plugin.hs new file mode 100644 --- /dev/null +++ b/ECTA/Plugin.hs @@ -0,0 +1,149 @@ +{-# 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) +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) + + +plugin :: Plugin +plugin = + defaultPlugin + { holeFitPlugin = \opts -> + Just $ + HoleFitPluginR { + hfPluginInit = newTcRef [], + hfPluginRun = \ref -> + ( HoleFitPlugin + { candPlugin = \_ c -> writeTcRef ref c >> return [], + fitPlugin = \h _ -> readTcRef ref >>= ectaPlugin opts h + } + ), + 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 [] = [] + +ectaPlugin :: [CommandLineOption] -> TypedHole -> [HoleFitCandidate] -> TcM [HoleFit] +ectaPlugin opts TyH{..} scope | Just hole <- tyHCt, + ty <- ctPred hole = 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 + 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 + hsc_env <- getTopEnv + instance_comps <- mapMaybe instToTerm . concat <$> + mapMaybeM (fmap (fmap (\(_,_,c,_,_) -> c) . snd) + . liftIO . tcRnGetInfo hsc_env . getName + . tyConAppTyCon) constraints + let 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) = + let argNodes = ngnodes local_scope_comps + addSyms st tt = map (Bi.bimap (Symbol . st) (tt . typeToFta)) + 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) + case typeToSkeleton ty of + Just (t, cons) | resNode <- typeToFta t -> do + let res = getAllTerms $ refold $ reduceFully $ filterType scopeNode resNode + ppterms <- concatMapM (prettyMatch skels groups . prettyTerm ) res + let even_more_terms = + map (pp . prettyTerm) $ + concatMap (getAllTerms . refold . reduceFully . flip filterType resNode ) + (rtkUpToKAtLeast1 argNodes scope_comps anyArg True 7) + return $ map (RawHoleFit . text . unpack) $ ppterms ++ even_more_terms + _ -> do liftIO $ putStrLn $ "Could not skeleton `" ++ showSDocUnsafe (ppr ty) ++"`" + return [] + + +-- TODO: +-- 1. I think we need to add type applications, i.e. [] @a or similar, let's see. diff --git a/ECTA/Plugin/Utils.hs b/ECTA/Plugin/Utils.hs new file mode 100644 --- /dev/null +++ b/ECTA/Plugin/Utils.hs @@ -0,0 +1,304 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE OverloadedStrings #-} +module ECTA.Plugin.Utils where + +-- These will be ruined by GHC 9.0+ due to the reorganization +-- These will be ruined by GHC 9.0+ due to the reorganization +import GhcPlugins hiding ((<>)) +import TcRnTypes + +import Application.TermSearch.Dataset +import Application.TermSearch.Type +import GhcPlugins hiding ((<>)) +import Type +import Data.Text (pack, Text, unpack) +import Data.Maybe (mapMaybe, isJust, fromJust) + +import Data.Map (Map) +import qualified Data.Map as Map + +import GHC.IO.Unsafe +import Data.IORef +import TcRnMonad (newUnique, getTopEnv, getLocalRdrEnv, getGlobalRdrEnv) +import TcEnv (tcLookupTyCon) +import Data.Char (isAlpha, isAscii) +import Data.ECTA (Node (Node), mkEdge, Edge (Edge), pathsMatching, mapNodes, createMu) +import Data.ECTA.Term +import Application.TermSearch.Utils (theArrowNode, arrowType, var1, var2, var3, var4, varAcc, mkDatatype) +import Data.ECTA (union) +import Data.ECTA.Paths (getPath, mkEqConstraints, path, Path) +import Debug.Trace (traceShow) +import qualified Data.Monoid as M +import Data.List (groupBy, sortOn, permutations) +import Data.Function (on) +import Data.Tuple (swap) +import Data.Containers.ListUtils (nubOrd) +import Data.List (subsequences) + + -- The old "global variable" trick, as we are creating new type variables + -- from scratch, but we want all 'a's to refer to the same variabel, etc. +tyVarMap :: IORef (Map Text TyVar) +{-# NOINLINE tyVarMap #-} +tyVarMap = unsafePerformIO (newIORef Map.empty) + +skeletonToType :: TypeSkeleton -> TcM (Maybe Type) +skeletonToType (TVar t) = Just <$> + do tm <- liftIO $ readIORef tyVarMap + case tm Map.!? t of + Just tv -> return $ mkTyVarTy tv + _ -> do + unq <- newUnique + let nm = mkSystemName unq $ mkOccName tvName (unpack t) + ntv = mkTyVar nm liftedTypeKind + liftIO $ modifyIORef tyVarMap (Map.insert t ntv) + return $ mkTyVarTy ntv +skeletonToType (TFun arg ret) = + do argty <- skeletonToType arg + retty <- skeletonToType ret + case (argty, retty) of + (Just a, Just r) -> return (Just (mkVisFunTy a r)) + _ -> return Nothing + -- This will be mkVisFunTyMany in 9.0+ +skeletonToType (TCons con sk) = + do args <- mapM skeletonToType sk + let occn = mkOccName tcName (unpack con) + lcl_env <- getLocalRdrEnv + gbl_env <- getGlobalRdrEnv + let conName = case lookupLocalRdrOcc lcl_env occn of + Just res -> Just res + _ -> case lookupGlobalRdrEnv gbl_env occn of + -- Note: does not account for shadowing + (GRE{..}:_) -> Just gre_name + _ -> Nothing + case conName of + Just n | all isJust args, argTys <- map fromJust args -> do + conTy <- tcLookupTyCon n + return $ Just $ mkTyConApp conTy argTys + _ -> return Nothing + + + + +-- | Extremely crude at the moment!! +-- Returns the typeSkeleton and any constructors (for instance lookup) +typeToSkeleton :: Type -> Maybe (TypeSkeleton, [Type]) +typeToSkeleton ty | (vars@(_:_), r) <- splitForAllTys ty, + all valid vars + = typeToSkeleton r + where + valid :: TyCoVar -> Bool + -- for now + valid = (`elem` ["a", "b", "c", "d", "acc"]) . showSDocUnsafe . ppr +typeToSkeleton ty | isTyVarTy ty, + Just tt <- tyVarToSkeletonText ty = Just (TVar tt, []) +typeToSkeleton ty | Just (arg, ret) <- splitFunTy_maybe ty, + Just (argsk, ac)<- typeToSkeleton arg, + Just (retsk, rc)<- typeToSkeleton ret = + Just (TFun argsk retsk,ac) +typeToSkeleton ty | (cons, args@(_:_)) <- splitAppTys ty, + Just const <- typeToSkeletonText cons, + argsks <- mapMaybe typeToSkeleton args, + (aks, acs) <- unzip argsks = + Just (TCons const aks, cons:concat acs) +typeToSkeleton ty | (cons, []) <- splitAppTys ty, + Just const <- typeToSkeletonText cons + -- These are the ones we don't handle so far + = Just (TCons const [], [cons]) + +-- TODO: Filter out lifted rep etc. +typeToSkeletonText :: Outputable a => a -> Maybe Text +typeToSkeletonText ty = Just $ pack $ showSDocUnsafe $ ppr ty + +-- TODO: make sure it's one of the supported ones +tyVarToSkeletonText :: Outputable a => a -> Maybe Text +tyVarToSkeletonText ty = Just $ pack $ stripNumbers $ showSDocUnsafe $ ppr ty + -- by default, the variables are given a number (e.g. a1). + -- we just strip them off for now. + where stripNumbers :: String -> String + stripNumbers = takeWhile isAlpha + + +constFunc :: Symbol -> Node -> Edge +constFunc s t = Edge s [t] + +applyOperator :: Comps -> Node +applyOperator comps = Node + [ constFunc + "$" + (generalize comps $ arrowType (arrowType var1 var2) (arrowType var1 var2)) + , constFunc "id" (generalize comps $ arrowType var1 var1) + ] + +tk :: Comps -> Node -> Bool -> Int -> [Node] +tk _ _ _ 0 = [] +tk _ anyArg False 1 = [anyArg] +tk comps anyArg True 1 = [anyArg, applyOperator comps] +tk comps anyArg _ k = map constructApp [1 .. (k-1)] + where + constructApp :: Int -> Node + constructApp i = + mapp (union $ tk comps anyArg False i) + (union $ tk comps anyArg True (k-i)) + +tkUpToK :: Comps -> Node -> Bool -> Int -> [Node] +tkUpToK comps anyArg includeApp k = concatMap (tk comps anyArg includeApp) [1..k] + +-- type Argument = (Symbol, Node) +rtk :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node] +rtk [] comps anyArg includeApplyOp k = tk comps anyArg False k +rtk [(s,t)] _ _ _ 1 = [Node [constFunc s t]] -- If we have one arg we use it +rtk args _ _ _ k | k < length args = [] +rtk args comps anyArg _ k = concatMap (\i -> map (constructApp i) allSplits) [1..(k-1)] + where allSplits = map (`splitAt` args) [0.. (length args)] + constructApp :: Int -> ([Argument], [Argument]) -> Node + constructApp i (xs, ys) = + let f = union $ rtk xs comps anyArg False i + x = union $ rtk ys comps anyArg True (k-i) + in mapp f x + +rtkOfSize :: [Argument] -> Comps -> Node -> Bool -> Int -> Node +rtkOfSize args comps anyArg includeApp k = + union $ concatMap (\a -> rtk a comps anyArg includeApp k) $ permutations args + +rtkUpToK :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node] +rtkUpToK args comps anyArg includeApp k = + map (rtkOfSize args comps anyArg includeApp) [1..k] + +rtkAtLeast1 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node] +rtkAtLeast1 args comps anyArg includeApp k = + map (\as -> rtkOfSize as comps anyArg includeApp k) $ map (:[]) args + +-- rtkUpToKAtLeast1 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node] +-- rtkUpToKAtLeast1 args comps anyArg includeApp k = +-- concatMap (\as -> rtkUpToK as comps anyArg includeApp k) $ map (:[]) args + +-- Slower for some reason? Probably wrong also because there's a lot more +-- repition, since we don't exclude the args etc etc. +-- TODO: improve by e.g. removing the ones already used from comps etc. +rtkUpToKAtLeast1 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node] +rtkUpToKAtLeast1 args comps anyArg includeApp k = + concatMap (rtkAtLeast1 args comps anyArg includeApp) [1..k] + + +mapp :: Node -> Node -> Node +mapp n1 n2 = Node [ + mkEdge "app" + [getPath (path [0,2]) n1, theArrowNode, n1, n2] + (mkEqConstraints [ + [path [1], path [2, 0, 0]] -- it's a function + , [path [3, 0], path [2, 0, 1]] + , [path [0], path [2,0,2]] + ] + )] + + +chunks :: Int -> [a] -> [[a]] +chunks _ [] = [] +chunks n xs = f:chunks n nxt + where (f,nxt) = splitAt n xs + +removeDicts :: Term -> Term +removeDicts t = cleanup $ maybe t id $ rd t + where rd (Term (Symbol t) args) | up@('<':'@':_) <- unpack t, + '>':'@':_ <- reverse up = Nothing + | args' <- mapMaybe rd args = + if null args' && t == "app" + then Nothing + else Just $ Term (Symbol t) args' + cleanup (Term (Symbol "app") [arg]) = cleanup arg + cleanup (Term (Symbol t) args) = Term (Symbol t) $ map cleanup args + + + +pp :: Term -> Text +pp = parIfReq . mconcat . pp' False . removeDicts + where + parIfReq :: Text -> Text + parIfReq t | (s:_) <- unpack t, + s /= '(', + not (isAlpha s) = "(" <> t <> ")" + | otherwise = t + + pp' :: Bool -> Term -> [Text] + pp' _ (Term (Symbol t) []) = [t] + pp' par (Term (Symbol "app") (arg:rest)) | res@(_:_) <- concatMap (pp' True) rest = + [rpar <> wparifreq <> " " <> mconcat (concatMap (pp' True) rest) <> lpar] + | otherwise = [wparifreq] + where parg = pp arg + (rpar,lpar) = if par then ("(", ")") else ("","") + wparifreq = if length (words $ unpack parg) > 1 + then "(" <> parg <> ")" else parg + +allConstructors :: Comps -> [(Text, Int)] +allConstructors comps = nubOrd (concatMap (getConstructors . snd) comps) + where + getConstructors :: TypeSkeleton -> [(Text, Int)] + getConstructors (TVar _ ) = [] + getConstructors (TFun t1 t2) = getConstructors t1 ++ getConstructors t2 + getConstructors (TCons nm ts) = + (nm, length ts) : concatMap getConstructors ts + +type Comps = [(Text,TypeSkeleton)] + +mtau :: Comps -> Node +mtau comps = createMu + (\n -> union + ( (arrowType n n:globalTyVars) + ++ map (Node . (: []) . constructorToEdge n) usedConstructors + ) + ) + where + constructorToEdge :: Node -> (Text, Int) -> Edge + constructorToEdge n (nm, arity) = Edge (Symbol nm) (replicate arity n) + + usedConstructors = allConstructors comps + +globalTyVars :: [Node] +globalTyVars = [var1, var2, var3, var4, varAcc] + +generalize :: Comps -> Node -> Node +generalize comps n@(Node [_]) = Node + [mkEdge s ns' (mkEqConstraints $ map pathsForVar vars)] + where + vars = globalTyVars + nWithVarsRemoved = mapNodes (\x -> if x `elem` vars then mtau comps else x) n + (Node [Edge s ns']) = nWithVarsRemoved + + pathsForVar :: Node -> [Path] + pathsForVar v = pathsMatching (== v) n +generalize _ n = n -- error $ "cannot generalize: " ++ show n + +invertMap :: Ord b => Map.Map a b -> Map.Map b [a] +invertMap = toMap . groupBy ((==) `on` fst) . sortOn fst . map swap . Map.toList + where toMap = Map.fromList . map (\((a,r):rs) -> (a,r:map snd rs)) + + +prettyMatch :: Map.Map Text TypeSkeleton -> Map.Map Text [Text] -> Term -> TcM [Text] +prettyMatch skels groups (Term (Symbol t) _) = + do ty <- skeletonToType tsk + let str = case ty of + Just t -> pack (" :: " ++ showSDocUnsafe (ppr t)) + _ -> pack (" :: " ++ show tsk) + return $ map (M.<> str) terms + where tsk = case skels Map.!? t of + Just r -> r + _ -> skels Map.! (pack $ tail $ unpack t) -- for generalization + terms = case groups Map.!? t of + Just r -> r + _ -> groups Map.! (pack $ tail $ unpack t) + + +mtypeToFta :: TypeSkeleton -> Node +mtypeToFta (TVar "a" ) = var1 +mtypeToFta (TVar "b" ) = var2 +mtypeToFta (TVar "c" ) = var3 +mtypeToFta (TVar "d" ) = var4 +mtypeToFta (TVar "acc") = varAcc +-- TODO: lift this restriction +mtypeToFta (TVar v) = + error + $ "Current implementation only supports function signatures with type variables a, b, c, d, and acc, but got " + ++ show v +mtypeToFta (TFun t1 t2 ) = arrowType (mtypeToFta t1) (mtypeToFta t2) +mtypeToFta (TCons "Fun" [t1, t2]) = arrowType (mtypeToFta t1) (mtypeToFta t2) +mtypeToFta (TCons s ts ) = mkDatatype s (map mtypeToFta ts) diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2022 Matthías Páll Gissurarson + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/ecta-plugin.cabal b/ecta-plugin.cabal new file mode 100644 --- /dev/null +++ b/ecta-plugin.cabal @@ -0,0 +1,28 @@ +name: ecta-plugin +cabal-version: 1.24 +build-type: Simple +version: 0.1.0 +author: Matthías Páll Gissurarson +maintainer: mpg@mpg.is +synopsis: Hole-Fit Synthesis using ECTAs +description: + + A hole-fit plugin for GHC that uses an ECTA to synthesize valid hole-fits. + See for more details on ECTAs! + +category: Compiler Plugin +license: MIT +license-file: LICENSE + +library + default-language: Haskell2010 + build-depends: base >= 4 && < 5, + ghc > 8.10 && < 9, + text > 1.2 && <= 1.3, + containers > 0.6 && <= 0.7, + ecta >= 1.0 && < 1.1 + exposed-modules: + ECTA.Plugin, + ECTA.Plugin.Utils + +