{-# 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 :: IORef (Map Text TyVar)
tyVarMap = IO (IORef (Map Text TyVar)) -> IORef (Map Text TyVar)
forall a. IO a -> a
unsafePerformIO (Map Text TyVar -> IO (IORef (Map Text TyVar))
forall a. a -> IO (IORef a)
newIORef Map Text TyVar
forall k a. Map k a
Map.empty)

skeletonToType :: TypeSkeleton -> TcM (Maybe Type)
skeletonToType :: TypeSkeleton -> TcM (Maybe Type)
skeletonToType (TVar Text
t) = Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Maybe Type)
-> IOEnv (Env TcGblEnv TcLclEnv) Type -> TcM (Maybe Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
     do Map Text TyVar
tm <- IO (Map Text TyVar)
-> IOEnv (Env TcGblEnv TcLclEnv) (Map Text TyVar)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Map Text TyVar)
 -> IOEnv (Env TcGblEnv TcLclEnv) (Map Text TyVar))
-> IO (Map Text TyVar)
-> IOEnv (Env TcGblEnv TcLclEnv) (Map Text TyVar)
forall a b. (a -> b) -> a -> b
$ IORef (Map Text TyVar) -> IO (Map Text TyVar)
forall a. IORef a -> IO a
readIORef IORef (Map Text TyVar)
tyVarMap
        case Map Text TyVar
tm Map Text TyVar -> Text -> Maybe TyVar
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? Text
t of
            Just TyVar
tv -> Type -> IOEnv (Env TcGblEnv TcLclEnv) Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> IOEnv (Env TcGblEnv TcLclEnv) Type)
-> Type -> IOEnv (Env TcGblEnv TcLclEnv) Type
forall a b. (a -> b) -> a -> b
$ TyVar -> Type
mkTyVarTy TyVar
tv
            Maybe TyVar
_ -> do
                Unique
unq <- TcRnIf TcGblEnv TcLclEnv Unique
forall gbl lcl. TcRnIf gbl lcl Unique
newUnique
                let nm :: Name
nm = Unique -> OccName -> Name
mkSystemName Unique
unq (OccName -> Name) -> OccName -> Name
forall a b. (a -> b) -> a -> b
$ NameSpace -> String -> OccName
mkOccName NameSpace
tvName (Text -> String
unpack Text
t)
                    ntv :: TyVar
ntv = Name -> Type -> TyVar
mkTyVar Name
nm Type
liftedTypeKind
                IO () -> IOEnv (Env TcGblEnv TcLclEnv) ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IOEnv (Env TcGblEnv TcLclEnv) ())
-> IO () -> IOEnv (Env TcGblEnv TcLclEnv) ()
forall a b. (a -> b) -> a -> b
$ IORef (Map Text TyVar)
-> (Map Text TyVar -> Map Text TyVar) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef (Map Text TyVar)
tyVarMap (Text -> TyVar -> Map Text TyVar -> Map Text TyVar
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
t TyVar
ntv)
                Type -> IOEnv (Env TcGblEnv TcLclEnv) Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> IOEnv (Env TcGblEnv TcLclEnv) Type)
-> Type -> IOEnv (Env TcGblEnv TcLclEnv) Type
forall a b. (a -> b) -> a -> b
$ TyVar -> Type
mkTyVarTy TyVar
ntv
skeletonToType (TFun TypeSkeleton
arg TypeSkeleton
ret) =
     do Maybe Type
argty <- TypeSkeleton -> TcM (Maybe Type)
skeletonToType TypeSkeleton
arg
        Maybe Type
retty <- TypeSkeleton -> TcM (Maybe Type)
skeletonToType TypeSkeleton
ret
        case (Maybe Type
argty, Maybe Type
retty) of
            (Just Type
a, Just Type
r) -> Maybe Type -> TcM (Maybe Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Type -> Type
mkVisFunTy Type
a Type
r))
            (Maybe Type, Maybe Type)
_ -> Maybe Type -> TcM (Maybe Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Type
forall a. Maybe a
Nothing
        -- This will be mkVisFunTyMany in 9.0+
skeletonToType (TCons Text
con [TypeSkeleton]
sk) =
    do [Maybe Type]
args <- (TypeSkeleton -> TcM (Maybe Type))
-> [TypeSkeleton] -> IOEnv (Env TcGblEnv TcLclEnv) [Maybe Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeSkeleton -> TcM (Maybe Type)
skeletonToType [TypeSkeleton]
sk
       let occn :: OccName
occn = NameSpace -> String -> OccName
mkOccName NameSpace
tcName (Text -> String
unpack Text
con)
       LocalRdrEnv
lcl_env <- RnM LocalRdrEnv
getLocalRdrEnv
       GlobalRdrEnv
gbl_env <- TcRn GlobalRdrEnv
getGlobalRdrEnv
       let conName :: Maybe Name
conName = case LocalRdrEnv -> OccName -> Maybe Name
lookupLocalRdrOcc LocalRdrEnv
lcl_env OccName
occn of
                        Just Name
res -> Name -> Maybe Name
forall a. a -> Maybe a
Just Name
res
                        Maybe Name
_ -> case GlobalRdrEnv -> OccName -> [GlobalRdrElt]
lookupGlobalRdrEnv GlobalRdrEnv
gbl_env OccName
occn of
                          -- Note: does not account for shadowing
                          (GRE{Bool
[ImportSpec]
Parent
Name
gre_name :: GlobalRdrElt -> Name
gre_par :: GlobalRdrElt -> Parent
gre_lcl :: GlobalRdrElt -> Bool
gre_imp :: GlobalRdrElt -> [ImportSpec]
gre_imp :: [ImportSpec]
gre_lcl :: Bool
gre_par :: Parent
gre_name :: Name
..}:[GlobalRdrElt]
_) -> Name -> Maybe Name
forall a. a -> Maybe a
Just Name
gre_name
                          [GlobalRdrElt]
_ -> Maybe Name
forall a. Maybe a
Nothing
       case Maybe Name
conName of
           Just Name
n | (Maybe Type -> Bool) -> [Maybe Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe Type -> Bool
forall a. Maybe a -> Bool
isJust [Maybe Type]
args, [Type]
argTys <- (Maybe Type -> Type) -> [Maybe Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Maybe Type -> Type
forall a. HasCallStack => Maybe a -> a
fromJust [Maybe Type]
args ->  do
               TyCon
conTy <- Name -> TcM TyCon
tcLookupTyCon Name
n
               Maybe Type -> TcM (Maybe Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Type -> TcM (Maybe Type)) -> Maybe Type -> TcM (Maybe Type)
forall a b. (a -> b) -> a -> b
$ Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Maybe Type) -> Type -> Maybe Type
forall a b. (a -> b) -> a -> b
$ TyCon -> [Type] -> Type
mkTyConApp TyCon
conTy [Type]
argTys
           Maybe Name
_ -> Maybe Type -> TcM (Maybe Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Type
forall a. Maybe a
Nothing




-- | Extremely crude at the moment!!
-- Returns the typeSkeleton and any constructors (for instance lookup)
typeToSkeleton :: Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton :: Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton Type
ty | (vars :: [TyVar]
vars@(TyVar
_:[TyVar]
_), Type
r) <- Type -> ([TyVar], Type)
splitForAllTys Type
ty,
                    (TyVar -> Bool) -> [TyVar] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TyVar -> Bool
valid [TyVar]
vars
                    = Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton Type
r
  where
      valid :: TyCoVar -> Bool
      -- for now
      valid :: TyVar -> Bool
valid = (String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String
"a", String
"b", String
"c", String
"d", String
"acc"]) (String -> Bool) -> (TyVar -> String) -> TyVar -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SDoc -> String
showSDocUnsafe (SDoc -> String) -> (TyVar -> SDoc) -> TyVar -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVar -> SDoc
forall a. Outputable a => a -> SDoc
ppr
typeToSkeleton Type
ty | Type -> Bool
isTyVarTy Type
ty,
                       Just Text
tt <- Type -> Maybe Text
forall a. Outputable a => a -> Maybe Text
tyVarToSkeletonText Type
ty = (TypeSkeleton, [Type]) -> Maybe (TypeSkeleton, [Type])
forall a. a -> Maybe a
Just  (Text -> TypeSkeleton
TVar Text
tt, [])
typeToSkeleton Type
ty | Just (Type
arg, Type
ret) <- Type -> Maybe (Type, Type)
splitFunTy_maybe Type
ty,
                       Just (TypeSkeleton
argsk, [Type]
ac)<- Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton Type
arg,
                       Just (TypeSkeleton
retsk, [Type]
rc)<- Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton Type
ret =
    (TypeSkeleton, [Type]) -> Maybe (TypeSkeleton, [Type])
forall a. a -> Maybe a
Just (TypeSkeleton -> TypeSkeleton -> TypeSkeleton
TFun TypeSkeleton
argsk TypeSkeleton
retsk,[Type]
ac)
typeToSkeleton Type
ty | (Type
cons, args :: [Type]
args@(Type
_:[Type]
_)) <- Type -> (Type, [Type])
splitAppTys Type
ty,
                       Just Text
const <- Type -> Maybe Text
forall a. Outputable a => a -> Maybe Text
typeToSkeletonText Type
cons,
                       [(TypeSkeleton, [Type])]
argsks <- (Type -> Maybe (TypeSkeleton, [Type]))
-> [Type] -> [(TypeSkeleton, [Type])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Type -> Maybe (TypeSkeleton, [Type])
typeToSkeleton [Type]
args,
                       ([TypeSkeleton]
aks, [[Type]]
acs) <- [(TypeSkeleton, [Type])] -> ([TypeSkeleton], [[Type]])
forall a b. [(a, b)] -> ([a], [b])
unzip [(TypeSkeleton, [Type])]
argsks =
    (TypeSkeleton, [Type]) -> Maybe (TypeSkeleton, [Type])
forall a. a -> Maybe a
Just (Text -> [TypeSkeleton] -> TypeSkeleton
TCons Text
const [TypeSkeleton]
aks, Type
consType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[[Type]] -> [Type]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Type]]
acs)
typeToSkeleton Type
ty | (Type
cons, []) <- Type -> (Type, [Type])
splitAppTys Type
ty,
                     Just Text
const <- Type -> Maybe Text
forall a. Outputable a => a -> Maybe Text
typeToSkeletonText Type
cons
    -- These are the ones we don't handle so far
    = (TypeSkeleton, [Type]) -> Maybe (TypeSkeleton, [Type])
forall a. a -> Maybe a
Just (Text -> [TypeSkeleton] -> TypeSkeleton
TCons Text
const [], [Type
cons])

-- TODO: Filter out lifted rep etc.
typeToSkeletonText :: Outputable a => a -> Maybe Text
typeToSkeletonText :: a -> Maybe Text
typeToSkeletonText a
ty = Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$ String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ SDoc -> String
showSDocUnsafe (SDoc -> String) -> SDoc -> String
forall a b. (a -> b) -> a -> b
$ a -> SDoc
forall a. Outputable a => a -> SDoc
ppr a
ty

-- TODO: make sure it's one of the supported ones
tyVarToSkeletonText :: Outputable a => a -> Maybe Text
tyVarToSkeletonText :: a -> Maybe Text
tyVarToSkeletonText a
ty = Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$ String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String -> String
stripNumbers (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ SDoc -> String
showSDocUnsafe (SDoc -> String) -> SDoc -> String
forall a b. (a -> b) -> a -> b
$ a -> SDoc
forall a. Outputable a => a -> SDoc
ppr a
ty
  -- by default, the variables are given a number (e.g. a1).
  -- we just strip them off for now.
  where stripNumbers :: String -> String
        stripNumbers :: String -> String
stripNumbers = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Char -> Bool
isAlpha


constFunc :: Symbol -> Node -> Edge
constFunc :: Symbol -> Node -> Edge
constFunc Symbol
s Node
t = Symbol -> [Node] -> Edge
Edge Symbol
s [Node
t]

applyOperator :: Comps -> Node
applyOperator :: Comps -> Node
applyOperator Comps
comps = [Edge] -> Node
Node
  [ Symbol -> Node -> Edge
constFunc
    Symbol
"$"
    (Comps -> Node -> Node
generalize Comps
comps (Node -> Node) -> Node -> Node
forall a b. (a -> b) -> a -> b
$ Node -> Node -> Node
arrowType (Node -> Node -> Node
arrowType Node
var1 Node
var2) (Node -> Node -> Node
arrowType Node
var1 Node
var2))
  , Symbol -> Node -> Edge
constFunc Symbol
"id" (Comps -> Node -> Node
generalize Comps
comps  (Node -> Node) -> Node -> Node
forall a b. (a -> b) -> a -> b
$ Node -> Node -> Node
arrowType Node
var1 Node
var1)
  ]

tk :: Comps -> Node -> Bool -> Int -> [Node]
tk :: Comps -> Node -> Bool -> Int -> [Node]
tk Comps
_ Node
_ Bool
_ Int
0 = []
tk Comps
_ Node
anyArg Bool
False Int
1 = [Node
anyArg]
tk Comps
comps Node
anyArg Bool
True Int
1 = [Node
anyArg, Comps -> Node
applyOperator Comps
comps]
tk Comps
comps Node
anyArg Bool
_ Int
k = (Int -> Node) -> [Int] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Node
constructApp [Int
1 .. (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)]
 where
   constructApp :: Int -> Node
   constructApp :: Int -> Node
constructApp Int
i =
      Node -> Node -> Node
mapp ([Node] -> Node
union ([Node] -> Node) -> [Node] -> Node
forall a b. (a -> b) -> a -> b
$ Comps -> Node -> Bool -> Int -> [Node]
tk Comps
comps Node
anyArg Bool
False Int
i)
           ([Node] -> Node
union ([Node] -> Node) -> [Node] -> Node
forall a b. (a -> b) -> a -> b
$ Comps -> Node -> Bool -> Int -> [Node]
tk Comps
comps Node
anyArg Bool
True (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i))

tkUpToK :: Comps -> Node -> Bool -> Int -> [Node]
tkUpToK :: Comps -> Node -> Bool -> Int -> [Node]
tkUpToK Comps
comps Node
anyArg Bool
includeApp Int
k = (Int -> [Node]) -> [Int] -> [Node]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Comps -> Node -> Bool -> Int -> [Node]
tk Comps
comps Node
anyArg Bool
includeApp) [Int
1..Int
k]

-- type Argument = (Symbol, Node)
rtk :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtk :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtk [] Comps
comps Node
anyArg Bool
includeApplyOp Int
k = Comps -> Node -> Bool -> Int -> [Node]
tk Comps
comps Node
anyArg Bool
False Int
k
rtk [(Symbol
s,Node
t)] Comps
_ Node
_ Bool
_ Int
1 = [[Edge] -> Node
Node [Symbol -> Node -> Edge
constFunc Symbol
s Node
t]] -- If we have one arg we use it
rtk [Argument]
args Comps
_ Node
_ Bool
_ Int
k | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Argument] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Argument]
args = []
rtk [Argument]
args Comps
comps Node
anyArg Bool
_ Int
k = (Int -> [Node]) -> [Int] -> [Node]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Int
i -> (([Argument], [Argument]) -> Node)
-> [([Argument], [Argument])] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> ([Argument], [Argument]) -> Node
constructApp Int
i) [([Argument], [Argument])]
allSplits) [Int
1..(Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)]
  where allSplits :: [([Argument], [Argument])]
allSplits = (Int -> ([Argument], [Argument]))
-> [Int] -> [([Argument], [Argument])]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [Argument] -> ([Argument], [Argument])
forall a. Int -> [a] -> ([a], [a])
`splitAt` [Argument]
args) [Int
0.. ([Argument] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Argument]
args)]
        constructApp :: Int -> ([Argument], [Argument]) -> Node
        constructApp :: Int -> ([Argument], [Argument]) -> Node
constructApp Int
i ([Argument]
xs, [Argument]
ys) =
          let f :: Node
f = [Node] -> Node
union ([Node] -> Node) -> [Node] -> Node
forall a b. (a -> b) -> a -> b
$ [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtk [Argument]
xs Comps
comps Node
anyArg Bool
False Int
i
              x :: Node
x = [Node] -> Node
union ([Node] -> Node) -> [Node] -> Node
forall a b. (a -> b) -> a -> b
$ [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtk [Argument]
ys Comps
comps Node
anyArg Bool
True (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
          in Node -> Node -> Node
mapp Node
f Node
x

rtkOfSize :: [Argument] -> Comps -> Node -> Bool -> Int -> Node
rtkOfSize :: [Argument] -> Comps -> Node -> Bool -> Int -> Node
rtkOfSize [Argument]
args Comps
comps Node
anyArg Bool
includeApp Int
k =
    [Node] -> Node
union ([Node] -> Node) -> [Node] -> Node
forall a b. (a -> b) -> a -> b
$ ([Argument] -> [Node]) -> [[Argument]] -> [Node]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\[Argument]
a -> [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtk [Argument]
a Comps
comps Node
anyArg Bool
includeApp Int
k) ([[Argument]] -> [Node]) -> [[Argument]] -> [Node]
forall a b. (a -> b) -> a -> b
$ [Argument] -> [[Argument]]
forall a. [a] -> [[a]]
permutations [Argument]
args

rtkUpToK :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkUpToK :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkUpToK [Argument]
args Comps
comps Node
anyArg Bool
includeApp Int
k = 
    (Int -> Node) -> [Int] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map ([Argument] -> Comps -> Node -> Bool -> Int -> Node
rtkOfSize [Argument]
args Comps
comps Node
anyArg Bool
includeApp) [Int
1..Int
k]

rtkAtLeast1 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkAtLeast1 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkAtLeast1 [Argument]
args Comps
comps Node
anyArg Bool
includeApp Int
k = 
    ([Argument] -> Node) -> [[Argument]] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (\[Argument]
as -> [Argument] -> Comps -> Node -> Bool -> Int -> Node
rtkOfSize [Argument]
as Comps
comps Node
anyArg Bool
includeApp Int
k) ([[Argument]] -> [Node]) -> [[Argument]] -> [Node]
forall a b. (a -> b) -> a -> b
$ (Argument -> [Argument]) -> [Argument] -> [[Argument]]
forall a b. (a -> b) -> [a] -> [b]
map (Argument -> [Argument] -> [Argument]
forall a. a -> [a] -> [a]
:[]) [Argument]
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 :: [Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkUpToKAtLeast1 [Argument]
args Comps
comps Node
anyArg Bool
includeApp Int
k = 
    (Int -> [Node]) -> [Int] -> [Node]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Argument] -> Comps -> Node -> Bool -> Int -> [Node]
rtkAtLeast1 [Argument]
args Comps
comps Node
anyArg Bool
includeApp) [Int
1..Int
k]


mapp :: Node -> Node -> Node
mapp :: Node -> Node -> Node
mapp Node
n1 Node
n2 = [Edge] -> Node
Node [
    Symbol -> [Node] -> EqConstraints -> Edge
mkEdge Symbol
"app"
    [Path -> Node -> Emptyable Node
forall t t'. Pathable t t' => Path -> t -> Emptyable t'
getPath ([Int] -> Path
path [Int
0,Int
2]) Node
n1, Node
theArrowNode, Node
n1, Node
n2]
    ([[Path]] -> EqConstraints
mkEqConstraints [
          [[Int] -> Path
path [Int
1], [Int] -> Path
path [Int
2, Int
0, Int
0]] -- it's a function
        , [[Int] -> Path
path [Int
3, Int
0], [Int] -> Path
path [Int
2, Int
0, Int
1]]
        , [[Int] -> Path
path [Int
0], [Int] -> Path
path [Int
2,Int
0,Int
2]]
        ]
    )]

        
chunks :: Int -> [a] -> [[a]]
chunks :: Int -> [a] -> [[a]]
chunks Int
_ [] = []
chunks Int
n [a]
xs = [a]
f[a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
:Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
n [a]
nxt
  where ([a]
f,[a]
nxt) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs

removeDicts :: Term -> Term
removeDicts :: Term -> Term
removeDicts Term
t = Term -> Term
cleanup (Term -> Term) -> Term -> Term
forall a b. (a -> b) -> a -> b
$ Term -> (Term -> Term) -> Maybe Term -> Term
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Term
t Term -> Term
forall a. a -> a
id (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Maybe Term
rd Term
t
 where rd :: Term -> Maybe Term
rd (Term (Symbol Text
t) [Term]
args) | up :: String
up@(Char
'<':Char
'@':String
_) <- Text -> String
unpack Text
t,
                                   Char
'>':Char
'@':String
_ <- String -> String
forall a. [a] -> [a]
reverse String
up = Maybe Term
forall a. Maybe a
Nothing
                                 | [Term]
args' <- (Term -> Maybe Term) -> [Term] -> [Term]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Term -> Maybe Term
rd [Term]
args =
                                   if [Term] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Term]
args' Bool -> Bool -> Bool
&& Text
t Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"app"
                                   then Maybe Term
forall a. Maybe a
Nothing
                                   else Term -> Maybe Term
forall a. a -> Maybe a
Just (Term -> Maybe Term) -> Term -> Maybe Term
forall a b. (a -> b) -> a -> b
$ Symbol -> [Term] -> Term
Term (Text -> Symbol
Symbol Text
t) [Term]
args'
       cleanup :: Term -> Term
cleanup (Term (Symbol Text
"app") [Term
arg]) = Term -> Term
cleanup Term
arg
       cleanup (Term (Symbol Text
t) [Term]
args) = Symbol -> [Term] -> Term
Term (Text -> Symbol
Symbol Text
t) ([Term] -> Term) -> [Term] -> Term
forall a b. (a -> b) -> a -> b
$ (Term -> Term) -> [Term] -> [Term]
forall a b. (a -> b) -> [a] -> [b]
map Term -> Term
cleanup [Term]
args

parIfReq :: Text -> Text
parIfReq :: Text -> Text
parIfReq Text
t | (Char
s:String
_) <- Text -> String
unpack Text
t,
            Char
s Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'(',
            Bool -> Bool
not (Char -> Bool
isAlpha Char
s) = Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
            | Bool
otherwise = Text
t

pp :: Term -> Text
pp :: Term -> Text
pp = Text -> Text
parIfReq (Text -> Text) -> (Term -> Text) -> Term -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Text
ppNoPar

ppNoPar :: Term -> Text
ppNoPar :: Term -> Text
ppNoPar = [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat ([Text] -> Text) -> (Term -> [Text]) -> Term -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Term -> [Text]
pp' Bool
False (Term -> [Text]) -> (Term -> Term) -> Term -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Term
removeDicts
 where
  pp' :: Bool -> Term -> [Text]
  pp' :: Bool -> Term -> [Text]
pp' Bool
_ (Term (Symbol Text
t) [])  = [Text
t]
  pp' Bool
par (Term (Symbol Text
"app") (Term
arg:[Term]
rest)) | res :: [Text]
res@(Text
_:[Text]
_) <- (Term -> [Text]) -> [Term] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Bool -> Term -> [Text]
pp' Bool
True) [Term]
rest =
      [Text
rpar Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
wparifreq Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat ((Term -> [Text]) -> [Term] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Bool -> Term -> [Text]
pp' Bool
True) [Term]
rest) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
lpar]
                                           | Bool
otherwise = [Text
wparifreq]
    where parg :: Text
parg = Term -> Text
pp Term
arg
          (Text
rpar,Text
lpar) = if Bool
par then (Text
"(", Text
")") else (Text
"",Text
"")
          wparifreq :: Text
wparifreq = if [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (String -> [String]
words (String -> [String]) -> String -> [String]
forall a b. (a -> b) -> a -> b
$ Text -> String
unpack Text
parg) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
                      then Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
parg Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")" else Text
parg

allConstructors :: Comps -> [(Text, Int)]
allConstructors :: Comps -> [(Text, Int)]
allConstructors Comps
comps = [(Text, Int)] -> [(Text, Int)]
forall a. Ord a => [a] -> [a]
nubOrd (((Text, TypeSkeleton) -> [(Text, Int)]) -> Comps -> [(Text, Int)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (TypeSkeleton -> [(Text, Int)]
getConstructors (TypeSkeleton -> [(Text, Int)])
-> ((Text, TypeSkeleton) -> TypeSkeleton)
-> (Text, TypeSkeleton)
-> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, TypeSkeleton) -> TypeSkeleton
forall a b. (a, b) -> b
snd) Comps
comps)
 where
  getConstructors :: TypeSkeleton -> [(Text, Int)]
  getConstructors :: TypeSkeleton -> [(Text, Int)]
getConstructors (TVar Text
_    ) = []
  getConstructors (TFun TypeSkeleton
t1 TypeSkeleton
t2) = TypeSkeleton -> [(Text, Int)]
getConstructors TypeSkeleton
t1 [(Text, Int)] -> [(Text, Int)] -> [(Text, Int)]
forall a. [a] -> [a] -> [a]
++ TypeSkeleton -> [(Text, Int)]
getConstructors TypeSkeleton
t2
  getConstructors (TCons Text
nm [TypeSkeleton]
ts) =
    (Text
nm, [TypeSkeleton] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeSkeleton]
ts) (Text, Int) -> [(Text, Int)] -> [(Text, Int)]
forall a. a -> [a] -> [a]
: (TypeSkeleton -> [(Text, Int)]) -> [TypeSkeleton] -> [(Text, Int)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap TypeSkeleton -> [(Text, Int)]
getConstructors [TypeSkeleton]
ts

type Comps = [(Text,TypeSkeleton)]

mtau :: Comps -> Node
mtau :: Comps -> Node
mtau Comps
comps = (Node -> Node) -> Node
createMu
  (\Node
n -> [Node] -> Node
union
    (  (Node -> Node -> Node
arrowType Node
n Node
nNode -> [Node] -> [Node]
forall a. a -> [a] -> [a]
:[Node]
globalTyVars)
    [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
++ ((Text, Int) -> Node) -> [(Text, Int)] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map ([Edge] -> Node
Node ([Edge] -> Node) -> ((Text, Int) -> [Edge]) -> (Text, Int) -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Edge -> [Edge] -> [Edge]
forall a. a -> [a] -> [a]
: []) (Edge -> [Edge]) -> ((Text, Int) -> Edge) -> (Text, Int) -> [Edge]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Node -> (Text, Int) -> Edge
constructorToEdge Node
n) [(Text, Int)]
usedConstructors
    )
  )
 where
  constructorToEdge :: Node -> (Text, Int) -> Edge
  constructorToEdge :: Node -> (Text, Int) -> Edge
constructorToEdge Node
n (Text
nm, Int
arity) = Symbol -> [Node] -> Edge
Edge (Text -> Symbol
Symbol Text
nm) (Int -> Node -> [Node]
forall a. Int -> a -> [a]
replicate Int
arity Node
n)

  usedConstructors :: [(Text, Int)]
usedConstructors = Comps -> [(Text, Int)]
allConstructors Comps
comps

globalTyVars :: [Node]
globalTyVars :: [Node]
globalTyVars = [Node
var1, Node
var2, Node
var3, Node
var4, Node
varAcc]

generalize :: Comps -> Node -> Node
generalize :: Comps -> Node -> Node
generalize Comps
comps n :: Node
n@(Node [Edge
_]) = [Edge] -> Node
Node
  [Symbol -> [Node] -> EqConstraints -> Edge
mkEdge Symbol
s [Node]
ns' ([[Path]] -> EqConstraints
mkEqConstraints ([[Path]] -> EqConstraints) -> [[Path]] -> EqConstraints
forall a b. (a -> b) -> a -> b
$ (Node -> [Path]) -> [Node] -> [[Path]]
forall a b. (a -> b) -> [a] -> [b]
map Node -> [Path]
pathsForVar [Node]
vars)]
 where
  vars :: [Node]
vars                = [Node]
globalTyVars
  nWithVarsRemoved :: Node
nWithVarsRemoved    = (Node -> Node) -> Node -> Node
mapNodes (\Node
x -> if Node
x Node -> [Node] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Node]
vars then Comps -> Node
mtau Comps
comps else Node
x) Node
n
  (Node [Edge Symbol
s [Node]
ns']) = Node
nWithVarsRemoved

  pathsForVar :: Node -> [Path]
  pathsForVar :: Node -> [Path]
pathsForVar Node
v = (Node -> Bool) -> Node -> [Path]
pathsMatching (Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
v) Node
n
generalize Comps
_ Node
n = Node
n -- error $ "cannot generalize: " ++ show n

invertMap :: Ord b => Map.Map a b -> Map.Map b [a]
invertMap :: Map a b -> Map b [a]
invertMap = [[(b, a)]] -> Map b [a]
forall b. [[(b, b)]] -> Map b [b]
toMap ([[(b, a)]] -> Map b [a])
-> (Map a b -> [[(b, a)]]) -> Map a b -> Map b [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a) -> (b, a) -> Bool) -> [(b, a)] -> [[(b, a)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (b -> b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (b -> b -> Bool) -> ((b, a) -> b) -> (b, a) -> (b, a) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (b, a) -> b
forall a b. (a, b) -> a
fst) ([(b, a)] -> [[(b, a)]])
-> (Map a b -> [(b, a)]) -> Map a b -> [[(b, a)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a) -> b) -> [(b, a)] -> [(b, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (b, a) -> b
forall a b. (a, b) -> a
fst ([(b, a)] -> [(b, a)])
-> (Map a b -> [(b, a)]) -> Map a b -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> (b, a)) -> [(a, b)] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> (b, a)
forall a b. (a, b) -> (b, a)
swap ([(a, b)] -> [(b, a)])
-> (Map a b -> [(a, b)]) -> Map a b -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
Map.toList
  where toMap :: [[(b, b)]] -> Map b [b]
toMap = [(b, [b])] -> Map b [b]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(b, [b])] -> Map b [b])
-> ([[(b, b)]] -> [(b, [b])]) -> [[(b, b)]] -> Map b [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(b, b)] -> (b, [b])) -> [[(b, b)]] -> [(b, [b])]
forall a b. (a -> b) -> [a] -> [b]
map (\((b
a,b
r):[(b, b)]
rs) -> (b
a,b
rb -> [b] -> [b]
forall a. a -> [a] -> [a]
:((b, b) -> b) -> [(b, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b, b) -> b
forall a b. (a, b) -> b
snd [(b, b)]
rs))


prettyMatch :: Map.Map Text TypeSkeleton -> Map.Map Text [Text] -> Term -> TcM [Text]
prettyMatch :: Map Text TypeSkeleton -> Map Text [Text] -> Term -> TcM [Text]
prettyMatch Map Text TypeSkeleton
skels Map Text [Text]
groups (Term (Symbol Text
t) [Term]
_) =
  do Maybe Type
ty <- TypeSkeleton -> TcM (Maybe Type)
skeletonToType TypeSkeleton
tsk
     let str :: Text
str = case Maybe Type
ty of
               Just Type
t  -> String -> Text
pack (String
" :: " String -> String -> String
forall a. [a] -> [a] -> [a]
++  SDoc -> String
showSDocUnsafe (Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr Type
t))
               Maybe Type
_ -> String -> Text
pack (String
" :: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeSkeleton -> String
forall a. Show a => a -> String
show TypeSkeleton
tsk)
     [Text] -> TcM [Text]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Text] -> TcM [Text]) -> [Text] -> TcM [Text]
forall a b. (a -> b) -> a -> b
$ (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
M.<> Text
str) [Text]
terms
  where tsk :: TypeSkeleton
tsk = case Map Text TypeSkeleton
skels Map Text TypeSkeleton -> Text -> Maybe TypeSkeleton
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? Text
t of
                Just TypeSkeleton
r -> TypeSkeleton
r
                Maybe TypeSkeleton
_ -> Map Text TypeSkeleton
skels Map Text TypeSkeleton -> Text -> TypeSkeleton
forall k a. Ord k => Map k a -> k -> a
Map.! (String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. [a] -> [a]
tail (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Text -> String
unpack Text
t) -- for generalization
        terms :: [Text]
terms = case Map Text [Text]
groups Map Text [Text] -> Text -> Maybe [Text]
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? Text
t of
                  Just [Text]
r -> [Text]
r
                  Maybe [Text]
_ -> Map Text [Text]
groups Map Text [Text] -> Text -> [Text]
forall k a. Ord k => Map k a -> k -> a
Map.! (String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. [a] -> [a]
tail (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Text -> String
unpack Text
t)

-- | isSafe checks if we can use the type. Sadly ecta doesn't support all
-- types as of yet, so we fall back to the regulur valid hole-fits.
isSafe :: TypeSkeleton -> Bool
isSafe :: TypeSkeleton -> Bool
isSafe (TVar Text
"a"  ) = Bool
True
isSafe (TVar Text
"b"  ) = Bool
True
isSafe (TVar Text
"c"  ) = Bool
True
isSafe (TVar Text
"d"  ) = Bool
True
isSafe (TVar Text
"acc") = Bool
True
-- We would like to remove this restriction, but ecta does not support it yet.
isSafe (TVar Text
v) = Bool
False
isSafe (TFun  TypeSkeleton
t1    TypeSkeleton
t2      ) = TypeSkeleton -> Bool
isSafe TypeSkeleton
t1 Bool -> Bool -> Bool
&& TypeSkeleton -> Bool
isSafe TypeSkeleton
t2
isSafe (TCons Text
"Fun" [TypeSkeleton
t1, TypeSkeleton
t2]) = TypeSkeleton -> Bool
isSafe TypeSkeleton
t1 Bool -> Bool -> Bool
&& TypeSkeleton -> Bool
isSafe TypeSkeleton
t2
isSafe (TCons Text
s     [TypeSkeleton]
ts      ) = (TypeSkeleton -> Bool) -> [TypeSkeleton] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeSkeleton -> Bool
isSafe [TypeSkeleton]
ts