if-instance (empty) → 0.1.0.0

raw patch · 7 files changed

+651/−0 lines, 7 filesdep +basedep +ghcdep +ghc-tcplugin-api

Dependencies added: base, ghc, ghc-tcplugin-api, if-instance

Files

changelog.md +3/−0
example/M1.hs +25/−0
example/M2.hs +32/−0
if-instance.cabal +137/−0
src/Data/Constraint/If.hs +125/−0
src/IfCt/Plugin.hs +216/−0
test/Main.hs +113/−0

+ changelog.md view

@@ -0,0 +1,3 @@+# Version 0.1.0.0 (2021-08-30)
+
+Initial release.

+ example/M1.hs view

@@ -0,0 +1,25 @@+
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}
+
+module M1 where
+
+-- base
+import Data.Kind
+  ( Type )
+
+-- if-instance
+import Data.Constraint.If
+  ( IfCt(ifCt) )
+
+--------------------------------------------------------------------------------
+
+showFun :: forall (a :: Type). IfCt ( Show ( a -> a ) ) => ( a -> a ) -> String
+showFun = ifCt @( Show (a -> a) ) show ( \ _ -> "<<function>>" )
+
+test1 :: ( Bool -> Bool ) -> String
+test1 fun = showFun fun

+ example/M2.hs view

@@ -0,0 +1,32 @@+
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module M2 where
+
+import M1
+
+--------------------------------------------------------------------------------
+
+instance Show ( Bool -> Bool ) where
+  show f = show [ f False, f True ]
+
+test2 :: ( a -> a ) -> String
+test2 fun = showFun fun
+
+test3 :: ( Bool -> Bool ) -> String
+test3 fun = showFun fun
+
+test :: String
+test =
+  unlines
+    [ test1 not
+    , test2 not
+    , test3 not
+    , showFun not
+    ]

+ if-instance.cabal view

@@ -0,0 +1,137 @@+cabal-version:  3.0
+name:           if-instance
+version:        0.1.0.0
+synopsis:       Branch on whether a constraint is satisfied
+license:        BSD-3-Clause
+build-type:     Simple
+author:         Sam Derbyshire
+maintainer:     Sam Derbyshire
+copyright:      2021 Sam Derbyshire
+homepage:       https://github.com/sheaf/if-instance
+category:       Type System, Plugin
+description:
+
+  This library provides a mechanism that can be used to branch on
+  whether a constraint is satisfied (not limited to typeclass instances,
+  despite the name of the library).
+
+  Usage example:
+
+  @
+
+  {-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}
+
+  module MyModule where
+
+  import Data.Constraint.If ( IfCt(ifCt) )
+
+  hypot :: forall a. ( Floating a, IfCt (FMA a) ) => a -> a -> a
+  hypot = ifCt @(FMA a) withFMA withoutFMA
+    where
+      withFMA :: FMA a => a -> a -> a
+      withFMA a b =
+        let
+          h = sqrt $ fma a a (b * b)
+          h² = h * h
+          a² = a * a
+          x = fma (-b) b (h² - a²) + fma h h (-h²) - fma a a (-a²)
+        in
+          h - x / ( 2 * h )
+      withoutFMA :: a -> a -> a
+      withoutFMA a b = sqrt ( a * a + b * b )
+  @
+
+  Here we select between two ways of computing the hypotenuse function
+  based on whether we have access to the fused multiply-add operation
+
+  @ fma :: FMA a => a -> a -> a -> a @
+
+  which computes @ \\ a b c -> ( a * b ) + c @ in a single instruction,
+  providing stronger guarantees about precision of the resul.
+
+  A call of the form @hypot \@MyNumberType@ will either use the robust @withFMA@
+  function when an @FMA MyNumberType@ instance is available, or will fallback
+  to the simple @withoutFMA@ implementation when no such instance can be found.
+
+extra-source-files:
+  changelog.md
+
+common common
+
+  build-depends:
+    base
+      >= 4.14.0  && < 4.18,
+    ghc
+      >= 8.10    && < 9.6,
+
+  default-language:
+    Haskell2010
+
+  ghc-options:
+    -Wall
+    -Wcompat
+    -fwarn-missing-local-signatures
+    -fwarn-incomplete-uni-patterns
+    -fwarn-missing-deriving-strategies
+    -fno-warn-unticked-promoted-constructors
+
+library
+
+  import:
+    common
+
+  hs-source-dirs:
+    src
+
+  build-depends:
+    ghc-tcplugin-api
+      >= 0.5.0.0 && < 0.6,
+
+  exposed-modules:
+    Data.Constraint.If
+    IfCt.Plugin
+
+  default-language:
+    Haskell2010
+
+  ghc-options:
+    -Wall
+    -Wcompat
+    -fwarn-missing-local-signatures
+    -fwarn-incomplete-uni-patterns
+    -fwarn-missing-deriving-strategies
+    -fno-warn-unticked-promoted-constructors
+
+library if-instance-example
+
+  import:
+    common
+
+  hs-source-dirs:
+    example
+
+  build-depends:
+    if-instance
+
+  exposed-modules:
+    M2
+
+  other-modules:
+    M1
+
+test-suite if-instance-test
+
+  import:
+    common
+
+  type:
+    exitcode-stdio-1.0
+
+  build-depends:
+    if-instance
+
+  hs-source-dirs:
+    test
+
+  main-is:
+    Main.hs

+ src/Data/Constraint/If.hs view

@@ -0,0 +1,125 @@+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+
+{-|
+Module: Data.Constraint.If
+
+This module defines the typeclass 'IfCt', with method 'ifCt':
+
+> ifCt :: forall (ct :: Constraint) (r :: Type). IfCt ct => ( ct => r ) -> r -> r
+
+An expression of the form @ ifCt \@ct yes no @ denotes a selection between the two
+branches @yes@ and @no@:
+
+  - if the constraint @ct@ can be determined to hold at the point of solving @IfCt ct@,
+    then the @yes@ branch is selected, which has access to the @ct@ constraint;
+  - otherwise, the fallback branch @no@ is selected.
+
+To use this, you will also need to enable the corresponding 'IfCt.Plugin.plugin',
+by adding @\{\-\# OPTIONS_GHC -fplugin=IfCt.Plugin \#\-\}@
+to the header of your module.
+
+== Example
+
+We can select the more efficient 'nubOrd' function when an 'Ord' instance
+is available:
+
+> myNub :: forall (a :: Type). ( Eq a, IfCt (Ord a) ) => [a] -> [a]
+> myNub = ifCt @(Ord a) nubOrd nub
+>  -- 'nubOrd' when 'Ord a' is satisfied, 'nub' otherwise.
+
+When a user calls @myNub@, e.g.:
+
+> foo :: [(Int, Int)]
+> foo = myNub [(1,2), (3,3), (1,2), (2,2), (1,2), (1,4)]
+
+GHC will discharge the @IfCt (Ord (Int,Int))@ constraint by trying to solve
+the @Ord (Int, Int)@ constraint. In this case, GHC can solve the constraint
+using the two top-level instances (which we assume are in scope):
+
+> instance Ord Int
+> instance (Ord a, Ord b) => Ord (a,b)
+
+As the @ Ord (Int,Int) @ can be solved, GHC thus choose the first branch
+in 'ifCt', which in this case is 'nubOrd'.
+
+== When does branch selection occur?
+
+What is important to understand is that the branch selection happens
+precisely when the @IfCt ct@ constraint is solved.
+
+
+> { -# OPTIONS_GHC -fplugin=IfCt.Plugin #- }
+> module M1 where
+>
+> showFun :: forall (a :: Type). IfCt ( Show ( a -> a ) ) => ( a -> a ) -> String
+> showFun = ifCt @( Show (a -> a) ) show ( \ _ -> "<<function>>" )
+>
+> test1 :: ( Bool -> Bool ) -> String
+> test1 fun = showFun fun
+>
+> ----------------------------------------
+>
+> { -# OPTIONS_GHC -fplugin=IfCt.Plugin #- }
+> module M2 where
+>
+> import M1
+>
+> instance Show ( Bool -> Bool ) where
+>   show f = show [ f False, f True ]
+>
+> test2 :: ( a -> a ) -> String
+> test2 fun = showFun fun
+>
+> test3 :: ( Bool -> Bool ) -> String
+> test3 fun = showFun fun
+
+After loading @M2@, we get the following results:
+
+>>> test1 not
+"<<function>>"
+
+In this example, to typecheck @test1@ we need to solve @IfCt (Show (Bool -> Bool))@.
+As no instance for @Show (Bool -> Bool)@ is available in @M1@, we pick the second branch,
+resulting in @"\<\<function\>\>"@.
+
+>>> test2 not
+"<<function>>"
+
+In this example, we must solve @IfCt (Show (a -> a))@. There is no such instance in @M2@,
+so we pick the second branch.
+
+>>> test3 not
+"[True, False]"
+
+>>> showFun not
+"[True, False]"
+
+In these last two examples, we must solve @IfCt (Show (Bool -> Bool))@.
+Such an instance is in scope in @M2@, so we choose the first branch.
+-}
+
+module Data.Constraint.If
+  ( IfCt(..) )
+  where
+
+-- base
+import Data.Kind
+  ( Constraint )
+
+--------------------------------------------------------------------------------
+
+type IfCt :: Constraint -> Constraint
+class IfCt ct where
+  -- | @ ifCt \@ct a b@ returns @a@ if the constraint is satisfied,
+  -- and @b@ otherwise.
+  --
+  -- Requires the if-instance 'IfCt.Plugin.plugin':
+  -- add @{-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}@
+  -- to the header of your module.
+  --
+  -- Note: the selection happens at the point in the code where the @IfCt ct@
+  -- constraint is solved.
+  ifCt :: ( ct => r ) -> r -> r

+ src/IfCt/Plugin.hs view

@@ -0,0 +1,216 @@+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE PatternSynonyms #-}
+
+module IfCt.Plugin
+  ( plugin )
+  where
+
+-- base
+import Data.Maybe
+  ( catMaybes )
+#if !MIN_VERSION_ghc(9,2,0)
+import Unsafe.Coerce
+  ( unsafeCoerce )
+#endif
+
+-- ghc
+import GHC.Plugins
+  ( Plugin(..)
+  , defaultPlugin, purePlugin
+  )
+import GHC.Data.Bag
+  ( unitBag )
+import GHC.Tc.Solver.Interact
+  ( solveSimpleGivens, solveSimpleWanteds )
+import GHC.Tc.Solver.Monad
+  ( getTcEvBindsMap, readTcRef, runTcSWithEvBinds, traceTcS
+#if MIN_VERSION_ghc(9,2,0)
+  , wrapTcS
+#else
+  , TcS
+#endif
+  )
+import GHC.Tc.Types
+  ( TcM )
+import GHC.Utils.Outputable
+  ( (<+>), ($$), empty, text, vcat )
+
+-- ghc-tcplugin-api
+import GHC.TcPlugin.API
+import GHC.TcPlugin.API.Internal
+  ( unsafeLiftTcM )
+
+--------------------------------------------------------------------------------
+-- Plugin definition.
+
+-- | A type-checking plugin that solves @MyCt ct@ constraints.
+-- Theis allows users to branch on whether @ct@ is satisfied.
+--
+-- To use this plugin, add @{-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}@
+-- to your module header.
+--
+-- A @MyCt ct@ instance is solved by trying to solve @ct@:
+--
+--   - if solving succeeds, the 'Data.Constraint.If.ifCt' function will
+--     pick the first branch,
+--   - otherwise, 'Data.Constraint.If.ifCt' will pick the second branch.
+--
+-- This means that the branch selection occurs precisely at the moment
+-- at which we solve the @IfCt ct@ constraint.
+-- See the documentation of 'Data.Constraint.If.IfCt' for more information.
+plugin :: Plugin
+plugin =
+  defaultPlugin
+    { tcPlugin        = \ _args -> Just $ mkTcPlugin ifCtTcPlugin
+    , pluginRecompile = purePlugin
+    }
+
+ifCtTcPlugin :: TcPlugin
+ifCtTcPlugin =
+  TcPlugin
+    { tcPluginInit    = initPlugin
+    , tcPluginSolve   = solver
+    , tcPluginRewrite = \ _ -> emptyUFM
+    , tcPluginStop    = \ _ -> pure ()
+    }
+
+--------------------------------------------------------------------------------
+-- Plugin initialisation.
+
+data PluginDefs
+  = PluginDefs
+    { ifCtClass :: !Class }
+
+findModule :: MonadTcPlugin m => Maybe String -> String -> m Module
+findModule mb_pkg modName = do
+  findResult <- findImportedModule ( mkModuleName modName ) ( fmap fsLit mb_pkg )
+  case findResult of
+    Found _ res     -> pure res
+    FoundMultiple _ -> error $ "IfCt plugin: found multiple modules named " <> modName <> "."
+    _               -> error $ "IfCt plugin: could not find any module named " <> modName <> "."
+
+initPlugin :: TcPluginM Init PluginDefs
+initPlugin = do
+  ifCtModule <- findModule Nothing "Data.Constraint.If"
+  ifCtClass  <- tcLookupClass =<< lookupOrig ifCtModule ( mkClsOcc "IfCt" )
+  pure $ PluginDefs { ifCtClass }
+
+--------------------------------------------------------------------------------
+-- Constraint solving.
+
+solver :: PluginDefs -> [ Ct ] -> [ Ct ] -> TcPluginM Solve TcPluginSolveResult
+solver defs givens wanteds
+  | null wanteds
+  = pure $ TcPluginOk [] []
+  | otherwise
+  = do
+      tcPluginTrace "IfCt plugin {" (ppr givens $$ ppr wanteds)
+      solveds <- catMaybes <$> traverse ( solveWanted defs givens ) wanteds
+      tcPluginTrace "IfCt plugin }" empty
+      pure $ TcPluginOk solveds []
+
+solveWanted :: PluginDefs -> [ Ct ] -> Ct -> TcPluginM Solve ( Maybe ( EvTerm, Ct ) )
+solveWanted defs@( PluginDefs { ifCtClass } ) givens wanted
+  | ClassPred cls [ct_ty] <- classifyPredType ( ctPred wanted )
+  , cls == ifCtClass
+  = do
+    tcPluginTrace "IfCt plugin: found IfCt constraint" ( ppr wanted )
+    ct_ev <- newWanted ( ctLoc wanted ) ct_ty
+    let
+      ct :: Ct
+      ct = mkNonCanonical ct_ev
+      ct_ev_dest :: TcEvDest
+      ct_ev_dest = ctev_dest ct_ev
+    evBindsVar <- askEvBinds
+    -- Start a new Solver run.
+    unsafeLiftTcM $ runTcSWithEvBinds evBindsVar $ do
+      -- Add back all the Givens.
+      traceTcS "IfCt plugin: adding Givens to the inert set" (ppr givens)
+      solveSimpleGivens givens
+      -- Try to solve 'ct', using both Givens and top-level instances.
+      _ <- solveSimpleWanteds ( unitBag ct )
+      -- Now look up whether GHC has managed to produce evidence for 'ct'.
+      mb_ct_evTerm <-
+        case ct_ev_dest of
+          HoleDest ( CoercionHole { ch_ref = ref } ) -> do
+            mb_co <- readTcRef ref
+            traceTcS "IfCt plugin: coercion hole" (ppr mb_co)
+            case mb_co of
+              Nothing -> pure Nothing
+              Just co -> pure . Just $ evCoercion co
+          EvVarDest ev_var -> do
+            evBindsMap <- getTcEvBindsMap evBindsVar
+            let
+              mb_evBind :: Maybe EvBind
+              mb_evBind = lookupEvBind evBindsMap ev_var
+            traceTcS "IfCt plugin: evidence binding" (ppr mb_evBind)
+            case mb_evBind of
+              Nothing      -> pure Nothing
+              Just ev_bind -> pure . Just $ eb_rhs ev_bind
+      wanted_evTerm <- case mb_ct_evTerm of
+        Just ( EvExpr ct_evExpr ) -> do
+          -- We've managed to solve 'ct': use the evidence and take the 'True' branch.
+          traceTcS "IfCt plugin: constraint could be solved"
+            ( vcat
+              [ text "ct =" <+> ppr ct_ty
+              , text "ev =" <+> ppr ct_evExpr
+              ]
+            )
+          wrapTcS $ ifCtTrueEvTerm defs ct_ty ct_evExpr
+        _ -> do
+          -- We couldn't solve 'ct': take the 'False' branch.
+          traceTcS "IfCt plugin: constraint could not be solved"
+            ( text "ct =" <+> ppr ct_ty )
+          wrapTcS $ ifCtFalseEvTerm defs ct_ty
+      pure $ Just ( wanted_evTerm, wanted )
+  | otherwise
+  = pure Nothing
+
+-- Evidence term for @IfCt ct@ when @ct@ isn't satisfied.
+-- ifCt = \ @r (a :: ct => r) (_ :: r) -> a ct_evTerm
+ifCtTrueEvTerm :: PluginDefs -> Type -> EvExpr -> TcM EvTerm
+ifCtTrueEvTerm ( PluginDefs { ifCtClass } ) ct_ty ct_evTerm = do
+  r_name <- newName ( mkTyVarOcc "r" )
+  a_name <- newName ( mkVarOcc   "a" )
+  let
+    r, a, b :: CoreBndr
+    r = mkTyVar r_name liftedTypeKind
+    a = mkLocalId a_name Many ( mkInvisFunTyMany ct_ty r_ty )
+    b = mkWildValBinder Many r_ty
+    r_ty :: Type
+    r_ty = mkTyVarTy r
+  pure . EvExpr $
+    mkCoreConApps ( classDataCon ifCtClass )
+      [ Type ct_ty
+      , mkCoreLams [ r, a, b ]
+        ( mkCoreApps ( Var a ) [ ct_evTerm ] )
+      ]
+
+-- Evidence term for @IfCt ct@ when @ct@ isn't satisfied.
+-- ifCt = \ @r (_ :: ct => r) (b :: r) -> b
+ifCtFalseEvTerm :: PluginDefs -> Type -> TcM EvTerm
+ifCtFalseEvTerm ( PluginDefs { ifCtClass } ) ct_ty = do
+  r_name <- newName ( mkTyVarOcc "r" )
+  b_name <- newName ( mkVarOcc   "b" )
+  let
+    r, a, b :: CoreBndr
+    r = mkTyVar r_name liftedTypeKind
+    a = mkWildValBinder Many ( mkInvisFunTyMany ct_ty r_ty )
+    b = mkLocalId b_name Many r_ty
+    r_ty :: Type
+    r_ty = mkTyVarTy r
+  pure . EvExpr $
+    mkCoreConApps ( classDataCon ifCtClass )
+      [ Type ct_ty
+      , mkCoreLams [ r, a, b ] ( Var b )
+      ]
+
+--------------------------------------------------------------------------------
+
+#if !MIN_VERSION_ghc(9,2,0)
+wrapTcS :: TcM a -> TcS a
+wrapTcS = unsafeCoerce const
+#endif

+ test/Main.hs view

@@ -0,0 +1,113 @@+
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TypeApplications #-}
+
+{-# OPTIONS_GHC -fplugin=IfCt.Plugin #-}
+{-# OPTIONS_GHC -dcore-lint #-}
+
+module Main where
+
+-- base
+import Data.Maybe
+  ( mapMaybe )
+import Data.Kind
+  ( Constraint, Type )
+import System.Exit
+  ( exitFailure, exitSuccess )
+#if MIN_VERSION_ghc(9,4,0)
+import GHC.Exts
+  ( withDict )
+#endif
+
+-- IfCt
+import Data.Constraint.If
+  ( IfCt(ifCt) )
+
+--------------------------------------------------------------------------------
+
+type MyShow :: Type -> Constraint
+class MyShow a where
+  myShow :: a -> String
+
+instance MyShow Int where
+  myShow = show
+
+myShowAnything :: forall a. IfCt ( MyShow a ) => a -> String
+myShowAnything = ifCt @( MyShow a ) yes no
+  where
+    yes :: MyShow a => a -> String
+    yes = myShow
+    no :: a -> String
+    no _ = "<<no MyShow instance>>"
+
+-- Should use the "MyShow Int" instance.
+test1 :: String
+test1 = myShowAnything ( 123 :: Int )
+
+-- No "MyShow ( Int -> Int -> Int )" instance.
+test2 :: String
+test2 = myShowAnything ( (+) :: Int -> Int -> Int )
+
+data A = A
+
+myShowA :: IfCt ( MyShow A ) => String
+myShowA = myShowAnything A
+
+#if MIN_VERSION_ghc(9,4,0)
+-- Should use the instance locally provided by "withDict".
+test3 :: String
+test3 =
+  withDict @( A -> String ) @( MyShow A )
+    ( \ _ -> "A" )
+    myShowA
+#endif
+
+-- No "MyShow A" instance.
+test4 :: String
+test4 = myShowA
+
+--------------------------------------------------------------------------------
+
+data Test where
+  Test
+    :: ( Show a, Eq a )
+    => { testName     :: String
+       , testActual   :: a
+       , testExpected :: a
+       }
+    -> Test
+
+runTest :: Test -> Maybe String
+runTest ( Test { testName, testActual, testExpected } )
+  | testActual == testExpected
+  = Nothing
+  | otherwise
+  = Just $
+      "\n" <>
+      "Test '" <> testName <> "' failed.\n" <>
+      "Expected: " <> show testExpected <> "\n" <>
+      "  Actual: " <> show testActual
+
+tests :: [ Test ]
+tests =
+  [ Test "test1" test1 "123"
+  , Test "test2" test2 "<<no MyShow instance>>"
+#if MIN_VERSION_ghc(9,4,0)
+  , Test "test3" test3 "A"
+#endif
+  , Test "test4" test4 "<<no MyShow instance>>"
+  ]
+
+main :: IO ()
+main = do
+  let
+    results :: [ String ]
+    results = mapMaybe runTest tests
+  case results of
+    [] -> exitSuccess
+    _  -> putStrLn ( unlines results ) *> exitFailure