packages feed

ecta-plugin (empty) → 0.1.0

raw patch · 4 files changed

+503/−0 lines, 4 filesdep +basedep +containersdep +ecta

Dependencies added: base, containers, ecta, ghc, text

Files

+ ECTA/Plugin.hs view
@@ -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.
+ ECTA/Plugin/Utils.hs view
@@ -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)
+ LICENSE view
@@ -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.+
+ ecta-plugin.cabal view
@@ -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 <https://github.com/jkoppel/ecta> 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++