Plugin works + tests

polysemy-plugin/package.yaml

@@ -38,3 +38,20 @@ tests:
     - -with-rtsopts=-N
     dependencies:
     - polysemy-plugin
+    - hspec
+
+default-extensions:
+  - DataKinds
+  - DeriveFunctor
+  - FlexibleContexts
+  - GADTs
+  - LambdaCase
+  - PolyKinds
+  - RankNTypes
+  - ScopedTypeVariables
+  - StandaloneDeriving
+  - TypeApplications
+  - TypeOperators
+  - TypeFamilies
+
+  - UnicodeSyntax

polysemy-plugin/src/Polysemy/Plugin.hs

@@ -1,6 +1,93 @@
 {-# LANGUAGE NoMonomorphismRestriction #-}
 {-# LANGUAGE CPP #-}
 
+------------------------------------------------------------------------------
+-- The MIT License (MIT)
+--
+-- Copyright (c) 2017 Luka Horvat
+--
+-- 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.
+--
+------------------------------------------------------------------------------
+--
+-- This module is heavily based on 'Control.Effects.Plugin' from the
+-- 'simple-effects' package, originally by Luka Horvat.
+--
+-- https://gitlab.com/LukaHorvat/simple-effects/commit/966ce80b8b5777a4bd8f87ffd443f5fa80cc8845#f51c1641c95dfaa4827f641013f8017e8cd02aab
+
+
+------------------------------------------------------------------------------
+-- | A typechecker plugin that can disambiguate "obvious" uses of effects in
+-- Polysemy.
+--
+-- __Example:__
+--
+-- Consider the following program:
+--
+-- @
+-- foo :: 'Polysemy.Member' ('Polysemy.State.State' Int) r => 'Polysemy.Sem' r ()
+-- foo = 'Polysemy.State.put' 10
+-- @
+--
+-- What does this program do? Any human will tell you that it changes the state
+-- of the 'Int' to 10, which is clearly what's meant.
+--
+-- Unfortunately, Polysemy can't work this out on its own. Its reasoning is
+-- "maybe you wanted to change some other 'Polysemy.State.State' effect which
+-- is /also/ a 'Num', but you just forgot to add a 'Polysemy.Member' constraint
+-- for it."
+--
+-- This is obbviously insane, but it's the way the cookie crumbles.
+-- 'Polysemy.Plugin' is a typechecker plugin which will disambiguate the above
+-- program (and others) so the compiler will do what you want.
+--
+-- __Usage:__
+--
+-- Add the following line to your package configuration:
+--
+-- @
+-- ghc-options: -fplugin=Polysemy.Plugin
+-- @
+--
+-- __Limitations:__
+--
+-- The 'Polysemy.Plugin' will only disambiguate effects if there is exactly one
+-- relevant constraint in scope. For example, it will /not/ disambiguate the
+-- following program:
+--
+-- @
+-- bar :: 'Polysemy.Members' \'[ 'Polysemy.State.State' Int
+--                 , 'Polysemy.State.State' Double
+--                 ] r => 'Polysemy.Sem' r ()
+-- bar = 'Polysemy.State.put' 10
+-- @
+--
+-- because it is now unclear whether you're attempting to set the 'Int' or the
+-- 'Double'. Instead, you can manually write a type application in this case.
+--
+-- @
+-- bar :: 'Polysemy.Members' \'[ 'Polysemy.State.State' Int
+--                 , 'Polysemy.State.State' Double
+--                 ] r => 'Polysemy.Sem' r ()
+-- bar = 'Polysemy.State.put' @Int 10
+-- @
+--
 module Polysemy.Plugin
   ( plugin
   ) where
@@ -28,6 +115,7 @@ import TcSMonad hiding (tcLookupClass)
 import CoAxiom
 import Outputable
 
+
 plugin :: Plugin
 plugin = defaultPlugin
     { tcPlugin = const (Just fundepPlugin)
@@ -39,18 +127,19 @@ plugin = defaultPlugin
 fundepPlugin :: TcPlugin
 fundepPlugin = TcPlugin
     { tcPluginInit = do
-        md <- lookupModule (mkModuleName "Polysemy") (fsLit "polysemy")
-        monadEffectTcNm <- lookupName md (mkTcOcc "Member")
+        md <- lookupModule (mkModuleName "Polysemy.Internal.Union") (fsLit "polysemy")
+        monadEffectTcNm <- lookupName md (mkTcOcc "Find")
         tcLookupClass monadEffectTcNm
     , tcPluginSolve = solveFundep
     , tcPluginStop = const (return ()) }
 
 allMonadEffectConstraints :: Class -> [Ct] -> [(CtLoc, (Type, Type, Type))]
 allMonadEffectConstraints cls cts =
-    [ (ctLoc cd, (effName, eff, mon))
-        | cd@CDictCan{cc_class = cls', cc_tyargs = [eff, mon]} <- cts
+    [ (ctLoc cd, (effName, eff, r))
+        | cd@CDictCan{cc_class = cls', cc_tyargs = [_, r, eff]} <- cts
         , cls == cls'
-        , let (effName, _) = splitAppTys eff ]
+        , let effName = getEffName eff
+              ]
 
 singleListToJust :: [a] -> Maybe a
 singleListToJust [a] = Just a
@@ -63,17 +152,31 @@ findMatchingEffectIfSingular (effName, _, mon) ts = singleListToJust
         , eqType effName effName'
         , eqType mon mon' ]
 
+getEffName :: Type -> Type
+getEffName t = fst $ splitAppTys t
+
+
+mkWanted :: CtLoc -> Type -> Type -> TcPluginM (Maybe Ct)
+mkWanted loc eff eff' = do
+  if eqType (getEffName eff) (getEffName eff')
+     then do
+       (ev, _) <- unsafeTcPluginTcM $ runTcSDeriveds $ newWantedEq loc Nominal eff eff'
+       pure $ Just (CNonCanonical ev)
+     else
+       pure Nothing
+
+
 solveFundep :: Class -> [Ct] -> [Ct] -> [Ct] -> TcPluginM TcPluginResult
 solveFundep effCls giv _ want = do
-    pprPanic "wat" $ ppr $ giv ++ want
     let wantedEffs = allMonadEffectConstraints effCls want
     let givenEffs = snd <$> allMonadEffectConstraints effCls giv
-    eqs <- forM wantedEffs $ \(loc, e@(_, eff, _)) ->
-        case findMatchingEffectIfSingular e givenEffs of
-            Nothing -> return Nothing
-            Just eff' -> do
-                (ev, _) <- unsafeTcPluginTcM
-                    (runTcSDeriveds (newWantedEq loc Nominal eff eff'))
-                return (Just (CNonCanonical ev))
+    eqs <- forM wantedEffs $ \(loc, e@(_, eff, r)) ->
+      case findMatchingEffectIfSingular e givenEffs of
+        Nothing -> do
+          case splitAppTys r of
+            (_, [_, eff', _]) -> mkWanted loc eff eff'
+            _                 -> pure Nothing
+        Just eff' -> mkWanted loc eff eff'
+
     return (TcPluginOk [] (catMaybes eqs))
 

polysemy-plugin/test/Spec.hs

@@ -1,2 +1,134 @@
+{-# LANGUAGE AllowAmbiguousTypes        #-}
+{-# LANGUAGE BlockArguments             #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedStrings          #-}
+
+{-# OPTIONS_GHC -fplugin=Polysemy.Plugin #-}
+
+module Main where
+
+import Data.Functor.Identity
+import GHC.Exts
+import Polysemy
+import Polysemy.Error
+import Polysemy.State
+import Test.Hspec
+
+
+
+idState :: Member (State s) r => Sem r ()
+idState = do
+  s <- get
+  put s
+
+intState :: Member (State Int) r => Sem r ()
+intState = put 10
+
+numState :: Num a => Member (State a) r => Sem r ()
+numState = put 10
+
+strState :: Member (State String) r => Sem r ()
+strState = put "hello"
+
+oStrState :: IsString a => Member (State a) r => Sem r ()
+oStrState = put "hello"
+
+
+err :: Member (Error e) r => Sem r Bool
+err =
+  catch
+    do
+      throw undefined
+    \_ -> pure True
+
+
+errState :: Num s => Members '[Error e, State s] r => Sem r Bool
+errState = do
+  numState
+  err
+
+
+lifted :: Monad m => Member (Lift m) r => Sem r ()
+lifted = sendM $ pure ()
+
+
+newtype MyString = MyString String
+  deriving (IsString, Eq, Show)
+
+
 main :: IO ()
-main = putStrLn "Test suite not yet implemented"
+main = hspec $ do
+  describe "State effect" $ do
+    describe "get/put" $ do
+      it "should work in simple cases" $ do
+        flipShouldBe (True, ()) . run $ runState True idState
+
+      it "should, when polymorphic, eliminate the first matching effect" $ do
+        flipShouldBe (False, (True, ()))   . run $ runState False $ runState True idState
+
+      it "should, when polymorphic, not eliminate unmatching effects" $ do
+        flipShouldBe (True, Right @Int ()) . run $ runState True $ runError idState
+
+    describe "numbers" $ do
+      it "should interpret against concrete Int" $ do
+        flipShouldBe (10, ()) . run $ runState 0 intState
+
+      describe "polymorphic Num constraint" $ do
+        it "should interpret against Int" $ do
+          flipShouldBe (10 :: Int, ())     . run $ runState 0 numState
+
+        it "should interpret against Float" $ do
+          flipShouldBe (10 :: Float, ())   . run $ runState 0 numState
+
+        it "should interpret against Double" $ do
+          flipShouldBe (10 :: Double, ())  . run $ runState 0 numState
+
+        it "should interpret against Integer" $ do
+          flipShouldBe (10 :: Integer, ()) . run $ runState 0 numState
+
+    describe "strings" $ do
+      it "concrete interpret against concrete String" $ do
+        flipShouldBe ("hello", ()) . run $ runState "nothing" strState
+
+      describe "polymorphic IsString constraint" $ do
+        it "should interpret against String" $ do
+          flipShouldBe ("hello" :: String, ())   . run $ runState "nothing" oStrState
+
+        it "should interpret against MyString" $ do
+          flipShouldBe ("hello" :: MyString, ()) . run $ runState "nothing" oStrState
+
+
+  describe "Error effect" $ do
+    it "should interpret against Int" $ do
+      flipShouldBe (Right @Int True)  . run $ runError err
+    it "should interpret against Bool" $ do
+      flipShouldBe (Right @Bool True) . run $ runError err
+
+
+  describe "State/Error effect" $ do
+    it "should interpret against Int/String" $ do
+      flipShouldBe (10 :: Int, Right @String True)  . run $ runState 0 $ runError errState
+    it "should interpret against Float/Bool" $ do
+      flipShouldBe (10 :: Float, Right @Bool True)  . run $ runState 0 $ runError errState
+
+
+  describe "Error/State effect" $ do
+    it "should interpret against String/Int" $ do
+      flipShouldBe (Right @String (10 :: Int, True))  . run $ runError $ runState 0 errState
+    it "should interpret against Bool/Float" $ do
+      flipShouldBe (Right @Bool (10 :: Float, True))  . run $ runError $ runState 0 errState
+
+
+  describe "Lift effect" $ do
+    it "should interpret against IO" $ do
+      res <- runM lifted
+      res `shouldBe` ()
+
+    it "should interpret against Identity" $ do
+      let res = runM lifted
+      res `shouldBe` Identity ()
+
+
+flipShouldBe :: (Show a, Eq a) => a -> a -> Expectation
+flipShouldBe = flip shouldBe
+