Add defaulting plugins.

Like the built-in type defaulting rules these plugins can propose candidates
to resolve ambiguous type variables.

Machine learning and other large APIs like those for game engines introduce
new numeric types and other complex typed APIs. The built-in defaulting
mechanism isn't powerful enough to resolve ambiguous types in these cases forcing
users to specify minutia that they might not even know how to do. There is
an example defaulting plugin linked in the documentation. Applications include
defaulting the device a computation executes on, if a gradient should be
computed for a tensor, or the size of a tensor.

See https://github.com/ghc-proposals/ghc-proposals/pull/396 for details.

10 files changed, 262 insertions, 0 deletions

diff --git a/testsuite/tests/plugins/Makefile b/testsuite/tests/plugins/Makefile
index c00b26684b..6716aecbfb 100644
--- a/testsuite/tests/plugins/Makefile
+++ b/testsuite/tests/plugins/Makefile
@@ -163,3 +163,12 @@ T20218:
 T20218b:
 	"$(TEST_HC)" $(TEST_HC_OPTS) $(ghcPluginWayFlags) T20218b.hs -package-db simple-plugin/pkg.T20218b/local.package.conf -fplugin Simple.ReplacePlugin -package simple-plugin -v0
 	./T20218b
+
+.PHONY: test-defaulting-plugin
+test-defaulting-plugin:
+	"$(TEST_HC)" $(TEST_HC_OPTS) $(ghcPluginWayFlags) -v0 test-defaulting-plugin.hs -package-db defaulting-plugin/pkg.test-defaulting-plugin/local.package.conf
+
+.PHONY: test-defaulting-plugin-fail
+test-defaulting-plugin-fail:
+	-"$(TEST_HC)" $(TEST_HC_OPTS) $(ghcPluginWayFlags) -v0 test-defaulting-plugin-fail.hs -package-db defaulting-plugin/pkg.test-defaulting-plugin-fail/local.package.conf
+
diff --git a/testsuite/tests/plugins/all.T b/testsuite/tests/plugins/all.T
index 0fc41ec039..2146a2aa43 100644
--- a/testsuite/tests/plugins/all.T
+++ b/testsuite/tests/plugins/all.T
@@ -241,6 +241,7 @@ test('plugin-recomp-change-2',
       ],
      makefile_test, [])
 
+
 test('T20417',
      [extra_files(['plugin-recomp/', 'plugin-recomp-test.hs']),
       only_ways([config.ghc_plugin_way]),
@@ -248,6 +249,7 @@ test('T20417',
       ],
      makefile_test, [])
 
+
 test('T20218',
      [extra_files(['simple-plugin/']), only_ways([config.ghc_plugin_way]),
       pre_cmd('$MAKE -s --no-print-directory -C simple-plugin package.T20218 TOP={top}'),
@@ -261,3 +263,15 @@ test('T20218b',
       pre_cmd('$MAKE -s --no-print-directory -C simple-plugin package.T20218b TOP={top}')
       ],
      makefile_test, [])
+
+test('test-defaulting-plugin',
+     [extra_files(['defaulting-plugin/']),
+      only_ways([config.ghc_plugin_way]),
+      pre_cmd('$MAKE -s --no-print-directory -C defaulting-plugin package.test-defaulting-plugin TOP={top}')],
+     makefile_test, [])
+
+test('test-defaulting-plugin-fail',
+     [extra_files(['defaulting-plugin/']),
+      only_ways([config.ghc_plugin_way]),
+      pre_cmd('$MAKE -s --no-print-directory -C defaulting-plugin package.test-defaulting-plugin-fail TOP={top}')],
+     makefile_test, [])
diff --git a/testsuite/tests/plugins/defaulting-plugin/DefaultLifted.hs b/testsuite/tests/plugins/defaulting-plugin/DefaultLifted.hs
new file mode 100644
index 0000000000..62071b9cfb
--- /dev/null
+++ b/testsuite/tests/plugins/defaulting-plugin/DefaultLifted.hs
@@ -0,0 +1,112 @@
+{-# LANGUAGE MultiParamTypeClasses, KindSignatures, FlexibleInstances, DataKinds, PatternSynonyms, StandaloneDeriving, GeneralizedNewtypeDeriving, PolyKinds #-}
+{-# OPTIONS -Wno-orphans #-}
+module DefaultLifted(DefaultType,plugin) where
+import GHC.Plugins
+import GHC.Tc.Types.Constraint
+import GHC.Tc.Plugin
+import GHC.Core.InstEnv
+import GHC.Tc.Solver (approximateWC)
+import GHC.Unit.Finder (findPluginModule)
+import GHC.Driver.Config.Finder (initFinderOpts)
+import Data.List
+import GHC.Tc.Types
+import qualified Data.Map as M
+import Control.Monad (liftM2)
+import GHC.Tc.Utils.TcType
+
+class DefaultType x (y :: x)
+
+instance Eq Type where
+  (==) = eqType
+instance Ord Type where
+  compare = nonDetCmpType
+instance Semigroup (TcPluginM [a]) where
+  (<>) = liftM2 (++)
+instance Monoid (TcPluginM [a]) where
+  mempty = pure mempty
+
+plugin :: Plugin
+plugin = defaultPlugin {
+  defaultingPlugin = install,
+  pluginRecompile = purePlugin
+  }
+
+install :: p -> Maybe GHC.Tc.Types.DefaultingPlugin
+install _ = Just $ DefaultingPlugin { dePluginInit = initialize
+                                    , dePluginRun  = run
+                                    , dePluginStop = stop
+                                    }
+
+pattern FoundModule :: Module -> FindResult
+pattern FoundModule a <- Found _ a
+fr_mod :: a -> a
+fr_mod = id
+
+lookupModule :: ModuleName -- ^ Name of the module
+             -> TcPluginM Module
+lookupModule mod_nm = do
+  hsc_env <- getTopEnv
+  let dflags    = hsc_dflags hsc_env
+  let fopts     = initFinderOpts dflags
+  let fc        = hsc_FC hsc_env
+  let units     = hsc_units hsc_env
+  let home_unit = hsc_home_unit hsc_env
+  -- found_module <- findPluginModule fc fopts units home_unit mod_name
+  found_module <- tcPluginIO $ findPluginModule fc fopts units home_unit mod_nm
+  case found_module of
+    FoundModule h -> return (fr_mod h)
+    _          -> do
+      found_module' <- findImportedModule mod_nm $ Just $ fsLit "this"
+      case found_module' of
+        FoundModule h -> return (fr_mod h)
+        _          -> panicDoc "Unable to resolve module looked up by plugin: "
+                               (ppr mod_nm)
+
+data PluginState = PluginState { defaultClassName :: Name }
+
+-- | Find a 'Name' in a 'Module' given an 'OccName'
+lookupName :: Module -> OccName -> TcPluginM Name
+lookupName md occ = lookupOrig md occ
+
+solveDefaultType :: PluginState -> [Ct] -> TcPluginM DefaultingPluginResult
+solveDefaultType _     []      = return []
+solveDefaultType state wanteds = do
+  envs <- getInstEnvs
+  insts <- classInstances envs <$> tcLookupClass (defaultClassName state)
+  let defaults =
+        foldl' (\m inst ->
+                 case is_tys inst of
+                   [matchty, replacety] -> M.insertWith (++) matchty [replacety] m
+                   _ -> error "Unsupported defaulting type")
+               M.empty insts
+  let groups =
+        foldl' (\m wanted ->
+                  foldl' (\m' var -> M.insertWith (++) var [wanted] m')
+                         m
+                         (filter (isVariableDefaultable defaults) $ tyCoVarsOfCtList wanted))
+               M.empty wanteds
+  M.foldMapWithKey (\var cts ->
+                    case M.lookup (tyVarKind var) defaults of
+                      Nothing -> error "Bug, we already checked that this variable has a default"
+                      Just deftys -> do
+                        pure [DefaultingProposal var deftys cts])
+    groups
+  where isVariableDefaultable defaults v = isAmbiguousTyVar v && M.member (tyVarKind v) defaults
+
+lookupDefaultTypes :: TcPluginM PluginState
+lookupDefaultTypes = do
+    md   <- lookupModule (mkModuleName "DefaultLifted")
+    name <- lookupName md (mkTcOcc "DefaultType")
+    pure $ PluginState { defaultClassName = name }
+
+initialize :: TcPluginM PluginState
+initialize = do
+  lookupDefaultTypes
+
+run :: PluginState -> WantedConstraints -> TcPluginM DefaultingPluginResult
+run s ws = do
+  solveDefaultType s (ctsElts $ approximateWC False ws)
+
+stop :: Monad m => p -> m ()
+stop _ = do
+  return ()
diff --git a/testsuite/tests/plugins/defaulting-plugin/Makefile b/testsuite/tests/plugins/defaulting-plugin/Makefile
new file mode 100644
index 0000000000..7ce5b78e75
--- /dev/null
+++ b/testsuite/tests/plugins/defaulting-plugin/Makefile
@@ -0,0 +1,18 @@
+TOP=../../..
+include $(TOP)/mk/boilerplate.mk
+include $(TOP)/mk/test.mk
+
+clean.%:
+	rm -rf pkg.$*
+
+HERE := $(abspath .)
+$(eval $(call canonicalise,HERE))
+
+package.%:
+	$(MAKE) -s --no-print-directory clean.$*
+	mkdir pkg.$*
+	"$(TEST_HC)" -outputdir pkg.$* --make -v0 -o pkg.$*/setup Setup.hs
+	"$(GHC_PKG)" init pkg.$*/local.package.conf
+	pkg.$*/setup configure --distdir pkg.$*/dist -v0 $(CABAL_PLUGIN_BUILD) --prefix="$(HERE)/pkg.$*/install" --with-compiler="$(TEST_HC)" $(if $(findstring YES,$(HAVE_PROFILING)), --enable-library-profiling) --with-hc-pkg="$(GHC_PKG)" --package-db=pkg.$*/local.package.conf
+	pkg.$*/setup build     --distdir pkg.$*/dist -v0
+	pkg.$*/setup install   --distdir pkg.$*/dist -v0
diff --git a/testsuite/tests/plugins/defaulting-plugin/Setup.hs b/testsuite/tests/plugins/defaulting-plugin/Setup.hs
new file mode 100644
index 0000000000..e8ef27dbba
--- /dev/null
+++ b/testsuite/tests/plugins/defaulting-plugin/Setup.hs
@@ -0,0 +1,3 @@
+import Distribution.Simple
+
+main = defaultMain
diff --git a/testsuite/tests/plugins/defaulting-plugin/defaulting-plugin.cabal b/testsuite/tests/plugins/defaulting-plugin/defaulting-plugin.cabal
new file mode 100644
index 0000000000..a8f69ab7a0
--- /dev/null
+++ b/testsuite/tests/plugins/defaulting-plugin/defaulting-plugin.cabal
@@ -0,0 +1,10 @@
+name: defaulting-plugin
+cabal-version: >=1.24
+build-type: Simple
+version: 0.1.0.0
+
+library
+  default-language: Haskell2010
+  build-depends: base, ghc, containers
+  exposed-modules: DefaultLifted
+  ghc-options: -Wall
diff --git a/testsuite/tests/plugins/test-defaulting-plugin-fail.hs b/testsuite/tests/plugins/test-defaulting-plugin-fail.hs
new file mode 100644
index 0000000000..c751f3f5c1
--- /dev/null
+++ b/testsuite/tests/plugins/test-defaulting-plugin-fail.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE MultiParamTypeClasses, KindSignatures, FlexibleInstances, DataKinds, PolyKinds, RankNTypes, AllowAmbiguousTypes, TypeOperators, TypeFamilies, ScopedTypeVariables #-}
+{-# OPTIONS_GHC -fplugin DefaultLifted -fwarn-type-defaults #-}
+
+-- Tests defaulting plugins
+module Main where
+import GHC.TypeLits
+import Data.Proxy
+import DefaultLifted
+
+instance DefaultType Nat 4
+instance DefaultType Nat 2
+instance DefaultType Nat 0
+
+q :: forall (a :: Nat). (KnownNat a) => Integer
+q = natVal (Proxy :: Proxy a)
+
+w :: forall (a :: Nat). (KnownNat a, 2 <= a) => Integer
+w = natVal (Proxy :: Proxy a)
+
+e :: forall (a :: Nat). (KnownNat a, 5 <= a) => Integer
+e = natVal (Proxy :: Proxy a)
+
+main :: IO ()
+main = do
+  print $ q + w + e
diff --git a/testsuite/tests/plugins/test-defaulting-plugin-fail.stderr b/testsuite/tests/plugins/test-defaulting-plugin-fail.stderr
new file mode 100644
index 0000000000..c9843437b4
--- /dev/null
+++ b/testsuite/tests/plugins/test-defaulting-plugin-fail.stderr
@@ -0,0 +1,5 @@
+test-defaulting-plugin-fail.hs:25:19:
+     Cannot satisfy: 5 <= a0
+     In the second argument of ‘(+)’, namely ‘e’
+      In the second argument of ‘($)’, namely ‘q + w + e’
+      In a stmt of a 'do' block: print $ q + w + e
diff --git a/testsuite/tests/plugins/test-defaulting-plugin.hs b/testsuite/tests/plugins/test-defaulting-plugin.hs
new file mode 100644
index 0000000000..fce4e96a8a
--- /dev/null
+++ b/testsuite/tests/plugins/test-defaulting-plugin.hs
@@ -0,0 +1,29 @@
+{-# LANGUAGE MultiParamTypeClasses, KindSignatures, FlexibleInstances, DataKinds, PolyKinds, RankNTypes, AllowAmbiguousTypes, TypeOperators, TypeFamilies, ScopedTypeVariables #-}
+{-# OPTIONS_GHC -fplugin DefaultLifted -fwarn-type-defaults #-}
+
+-- Tests defaulting plugins
+module Main where
+import GHC.TypeLits
+import Data.Proxy
+import DefaultLifted
+
+instance DefaultType Nat 4
+instance DefaultType Nat 2
+instance DefaultType Nat 0
+
+class MyClass (a :: Nat) (b :: Nat) where
+  mc :: (KnownNat a, KnownNat b) => Proxy a -> Proxy b -> Integer
+
+instance MyClass a b where
+  mc a b = natVal a + natVal b
+
+q :: forall (a :: Nat). (KnownNat a) => Integer
+q = natVal (Proxy :: Proxy a)
+
+w :: forall (a :: Nat). (KnownNat a, 2 <= a) => Integer
+w = natVal (Proxy :: Proxy a)
+
+main :: IO ()
+main = do
+  print $ q + w
+  print $ mc Proxy Proxy
diff --git a/testsuite/tests/plugins/test-defaulting-plugin.stderr b/testsuite/tests/plugins/test-defaulting-plugin.stderr
new file mode 100644
index 0000000000..5108d5c805
--- /dev/null
+++ b/testsuite/tests/plugins/test-defaulting-plugin.stderr
@@ -0,0 +1,37 @@
+
+test-defaulting-plugin.hs:28:11: warning: [-Wtype-defaults (in -Wall)]
+     Defaulting the type variable ‘a0’ to type ‘0’ in the following constraint
+        KnownNat a0 arising from a use of ‘q’
+     In the first argument of ‘(+)’, namely ‘q’
+      In the second argument of ‘($)’, namely ‘q + w’
+      In a stmt of a 'do' block: print $ q + w
+
+test-defaulting-plugin.hs:28:15: warning: [-Wtype-defaults (in -Wall)]
+     Defaulting the type variable ‘a0’ to type ‘2’ in the following constraints
+        (KnownNat a0)
+          arising from a use of ‘w’ at test-defaulting-plugin.hs:28:15
+        (GHC.TypeError.Assert
+           (Data.Type.Ord.OrdCond (CmpNat 2 a0) 'True 'True 'False)
+           (TypeError ...))
+          arising from a use of ‘w’ at test-defaulting-plugin.hs:28:15
+     In the second argument of ‘(+)’, namely ‘w’
+      In the second argument of ‘($)’, namely ‘q + w’
+      In a stmt of a 'do' block: print $ q + w
+
+test-defaulting-plugin.hs:29:11: warning: [-Wtype-defaults (in -Wall)]
+     Defaulting the type variable ‘b0’ to type ‘0’ in the following constraint
+        KnownNat b0 arising from a use of ‘mc’
+     In the second argument of ‘($)’, namely ‘mc Proxy Proxy’
+      In a stmt of a 'do' block: print $ mc Proxy Proxy
+      In the expression:
+        do print $ q + w
+           print $ mc Proxy Proxy
+
+test-defaulting-plugin.hs:29:11: warning: [-Wtype-defaults (in -Wall)]
+     Defaulting the type variable ‘a0’ to type ‘0’ in the following constraint
+        KnownNat a0 arising from a use of ‘mc’
+     In the second argument of ‘($)’, namely ‘mc Proxy Proxy’
+      In a stmt of a 'do' block: print $ mc Proxy Proxy
+      In the expression:
+        do print $ q + w
+           print $ mc Proxy Proxy