Update smart constructors generation with tests and support for operators (#77)

src/Polysemy/Internal/TH/Effect.hs

@@ -36,15 +36,15 @@ import Prelude hiding ((<>))
 
 import Control.Monad
 import Data.Bifunctor
-import Data.Char                      (toLower)
+import Data.Char (toLower)
 import Data.Either
-import Data.Generics
+import Data.Generics hiding (Fixity)
 import Data.List
 import Data.Tuple
 import Language.Haskell.TH
 import Language.Haskell.TH.PprLib
 import Language.Haskell.TH.Datatype
-import Polysemy.Internal              (send, Member, Sem)
+import Polysemy.Internal (send, Member, Sem)
 import Polysemy.Internal.CustomErrors (DefiningModule)
 
 import qualified Data.Map.Strict as M
@@ -55,15 +55,18 @@ import qualified Data.Map.Strict as M
 -- | If @T@ is a GADT representing an effect algebra, as described in the
 -- module documentation for "Polysemy", @$('makeSem' ''T)@ automatically
 -- generates a smart constructor for every data constructor of @T@. This also
--- works for data family instances.
+-- works for data family instances. Names of smart constructors are created by
+-- changing first letter to lowercase or removing prefix @:@ in case of
+-- operators. Fixity declaration is preserved for both normal names and
+-- operators.
 --
 -- @since 0.1.2.0
 makeSem :: Name -> Q [Dec]
 makeSem = genFreer True
 
--- | Like 'makeSem', but does not provide type signatures. This can be used
+-- | Like 'makeSem', but does not provide type signatures and fixities. This
--- to attach Haddock comments to individual arguments for each generated
+-- can be used to attach Haddock comments to individual arguments for each
--- function.
+-- generated function.
 --
 -- @
 -- data Output o m a where
@@ -140,18 +143,20 @@ genFreer should_mk_sigs type_name = do
                          ]
   decs       <- traverse (genDec should_mk_sigs) cl_infos
 
-  let sigs = [ genSig <$> cl_infos | should_mk_sigs ]
+  let sigs = if should_mk_sigs then genSig <$> cl_infos else []
 
   return $ join $ def_mod_fi : sigs ++ decs
 
 ------------------------------------------------------------------------------
 -- | Generates signature for lifting function and type arguments to apply in
 -- its body on effect's data constructor.
-genSig :: ConLiftInfo -> Dec
+genSig :: ConLiftInfo -> [Dec]
 genSig cli
-  = SigD (cliFunName cli) $ quantifyType
+  =  maybe [] (pure . flip InfixD (cliFunName cli)) (cliFunFixity cli)
-  $ ForallT [] (member_cxt : cliFunCxt cli)
+  ++ [ SigD (cliFunName cli) $ quantifyType
-  $ foldArrows $ cliFunArgs cli ++ [sem `AppT` cliResType cli]
+       $ ForallT [] (member_cxt : cliFunCxt cli)
+       $ foldArrows $ cliFunArgs cli ++ [sem `AppT` cliResType cli]
+     ]
   where
     member_cxt = classPred ''Member [eff, VarT $ cliUnionName cli]
     eff        = foldl' AppT (ConT $ cliEffName cli) $ cliEffArgs cli
@@ -181,17 +186,27 @@ genDec should_mk_sigs cli = do
         ]
     ]
 
--------------------------------------------------------------------------------
+------------------------------------------------------------------------------
 -- | Info about constructor being lifted; use 'mkCLInfo' to create one.
 data ConLiftInfo = CLInfo
-  { cliEffName   :: Name    -- ^ name of effect's type constructor
+  { -- | Name of effect's type constructor
-  , cliEffArgs   :: [Type]  -- ^ effect specific type arguments
+    cliEffName   :: Name
-  , cliResType   :: Type    -- ^ result type specific to action
+    -- | Effect-specific type arguments
-  , cliConName   :: Name    -- ^ name of action constructor
+  , cliEffArgs   :: [Type]
-  , cliFunName   :: Name    -- ^ name of final function
+    -- | Result type specific to action
-  , cliFunArgs   :: [Type]  -- ^ final function arguments
+  , cliResType   :: Type
-  , cliFunCxt    :: Cxt     -- ^ constraints of final function
+    -- | Name of action constructor
-  , cliUnionName :: Name    -- ^ name of type variable parameterizing 'Sem'
+  , cliConName   :: Name
+    -- | Name of final function
+  , cliFunName   :: Name
+    -- | Fixity of function used as an operator
+  , cliFunFixity :: Maybe Fixity
+    -- | Final function arguments
+  , cliFunArgs   :: [Type]
+    -- | Constraints of final function
+  , cliFunCxt    :: Cxt
+    -- | Name of type variable parameterizing 'Sem'
+  , cliUnionName :: Name
   } deriving Show
 
 ------------------------------------------------------------------------------
@@ -204,7 +219,7 @@ mkCLInfo dti ci = do
   (raw_cli_eff_args, [m_arg, raw_cli_res_arg]) <-
     case splitAtEnd 2 $ datatypeInstTypes dti of
       r@(_, [_, _]) -> return r
-      _             -> missingEffTArgs cliEffName
+      _             -> missingEffArgs cliEffName
 
   m_name <-
     case tVarName m_arg of
@@ -212,6 +227,7 @@ mkCLInfo dti ci = do
       Nothing       -> mArgNotVar cliEffName m_arg
 
   cliUnionName <- newName "r"
+  cliFunFixity <- reifyFixity $ constructorName ci
 
   let normalizeType         = replaceMArg m_name cliUnionName
                             . simplifyKinds
@@ -229,8 +245,7 @@ mkCLInfo dti ci = do
       cliEffArgs            = normalizeType <$> raw_cli_eff_args
       cliResType            = normalizeType     raw_cli_res_arg
       cliConName            = constructorName ci
-      cliFunName            = mkName $ mapHead toLower $ nameBase
+      cliFunName            = liftFunNameFromCon cliConName
-                            $ constructorName ci
       cliFunArgs            = normalizeType <$> constructorFields ci
 
   return CLInfo{..}
@@ -243,8 +258,8 @@ mArgNotVar name mArg = fail $ show
   $  text "Monad argument ‘" <> ppr mArg <> text "’ in effect ‘"
   <> ppr name <> text "’ is not a type variable"
 
-missingEffTArgs :: Name -> Q a
+missingEffArgs :: Name -> Q a
-missingEffTArgs name = fail $ show
+missingEffArgs name = fail $ show
   $   text "Effect ‘" <> ppr name
       <> text "’ has not enough type arguments"
   $+$ nest 4
@@ -256,7 +271,7 @@ missingEffTArgs name = fail $ show
           )
       )
   where
-    base = mkName $ nameBase $ name
+    base = capturableBase name
     args = PlainTV . mkName <$> ["m", "a"]
 
 checkExtensions :: [Extension] -> Q ()
@@ -266,7 +281,7 @@ checkExtensions exts = do
         (\(ext, _) -> fail $ show
           $ char '‘' <> text (show ext) <> char '’'
             <+> text "extension needs to be enabled\
-                     \for smart constructors to work")
+                     \ for smart constructors to work")
         (find (not . snd) states)
 
 ------------------------------------------------------------------------------
@@ -281,6 +296,11 @@ capturableTVars = everywhere $ mkT $ \case
       goBndr (KindedTV n k) = KindedTV (capturableBase n) $ capturableTVars k
   t -> t
 
+------------------------------------------------------------------------------
+-- | Constructs capturable name from base of input name.
+capturableBase :: Name -> Name
+capturableBase = mkName . nameBase
+
 ------------------------------------------------------------------------------
 -- | Replaces use of @m@ in type with @Sem r@.
 replaceMArg :: TypeSubstitution t => Name -> Name -> t -> t
@@ -308,6 +328,15 @@ eqPairOrCxt p = case asEqualPred p of
   Just (a, VarT n) -> Left (n, a)
   _                -> Right p
 
+------------------------------------------------------------------------------
+-- | Creates name of lifting function from action name.
+liftFunNameFromCon :: Name -> Name
+liftFunNameFromCon n = mkName $ case nameBase n of
+                         ':':cs -> cs
+                         c  :cs -> toLower c : cs
+                         ""     -> error
+                           "liftFunNameFromCon: empty constructor name"
+
 ------------------------------------------------------------------------------
 -- | Folds a list of 'Type's into a right-associative arrow 'Type'.
 foldArrows :: [Type] -> Type
@@ -324,18 +353,7 @@ tVarName = \case
   ParensT t     -> tVarName t
   _             -> Nothing
 
-------------------------------------------------------------------------------
--- | Constructs capturable name from base of input name.
-capturableBase :: Name -> Name
-capturableBase = mkName . nameBase
-
 ------------------------------------------------------------------------------
 -- | 'splitAt' counting from the end.
 splitAtEnd :: Int -> [a] -> ([a], [a])
 splitAtEnd n = swap . join bimap reverse . splitAt n . reverse
-
-------------------------------------------------------------------------------
--- | Applies function to head of list, if there is one.
-mapHead :: (a -> a) -> [a] -> [a]
-mapHead _ []     = []
-mapHead f (x:xs) = f x : xs

test/ThEffectSpec.hs

@@ -0,0 +1,153 @@
+{-# LANGUAGE TemplateHaskell, GADTs, DataKinds, AllowAmbiguousTypes #-}
+
+module ThEffectSpec where
+
+import Polysemy
+import Test.Hspec
+import GHC.TypeLits
+import Data.Kind
+import Language.Haskell.TH hiding (Type)
+
+spec :: Spec
+spec = it "should compile" True
+
+-- Infix effects and actions -------------------------------------------------
+
+data (:#) m a where
+  (:#) :: a -> b -> m :# a
+
+infixl 4 :#
+
+makeSem ''(:#)
+
+reifyFixity '(#) >>= \case
+  Just (Fixity 4 InfixL) -> return []
+  _                      -> fail "Wrong fixity of generated operator"
+
+-- ADTs and ADT syntax -------------------------------------------------------
+
+data SimpleADT m a = SimpleADTC1 Int | SimpleADTC2 String
+
+makeSem ''SimpleADT
+
+data GADTSyntax m a where
+  GADTSyntaxC1 :: Int -> GADTSyntax m a
+  GADTSyntaxC2 :: String -> GADTSyntax m a
+
+makeSem ''GADTSyntax
+
+data ADTSyntax1 m a = (a ~ Int) => ADTSyntax1C String
+
+makeSem ''ADTSyntax1
+
+data ADTSyntax2 m a
+  = a ~ Int => ADTSyntax2C1 Int
+  | a ~ String => ADTSyntax2C2 String
+
+makeSem ''ADTSyntax2
+
+data ADTSyntax3 m a = Show a => ADTSyntax3C a
+
+makeSem ''ADTSyntax3
+
+-- We don't care about named fields (except that we accept them as names from
+-- effect in 'makeSem')
+data Fields m a = FieldsC { fieldsCF1 :: Int, fieldsCF2 :: String }
+
+makeSem ''Fields
+
+-- Newtypes ------------------------------------------------------------------
+
+newtype Newtype1 m a = Newtype1C Int
+
+makeSem ''Newtype1
+
+newtype Newtype2 m a where
+  Newtype2C :: String -> Newtype2 m a
+
+makeSem ''Newtype2
+
+-- Data families -------------------------------------------------------------
+
+data Instance = ADTI | GADTI | NTI
+
+data family Family (s :: Instance) (m :: Type -> Type) a
+
+data instance Family 'ADTI m a = ADTIC1 Int | ADTIC2 String
+
+makeSem 'ADTIC1
+
+data instance Family 'GADTI m a where
+  GADTIC1 :: Int -> Family 'GADTI m Int
+  GADTIC2 :: String -> Family 'GADTI m String
+
+makeSem 'GADTIC1
+
+newtype instance Family 'NTI m a = NTIC Int
+
+makeSem 'NTIC
+
+-- Phantom types -------------------------------------------------------------
+
+data Phantom m a
+
+makeSem ''Phantom
+
+-- Complex action types ------------------------------------------------------
+
+-- Inspired by:
+-- github.com/lexi-lambda/freer-simple/blob/ec84ae4e23ccba1ae05368100da750c196bbbcbb/tests/Tests/TH.hs#L37
+data Complex m a where
+  Mono            :: Int -> Complex m Bool
+  Poly            :: a -> Complex m a
+  PolyIn          :: a -> Complex m Bool
+  PolyOut         :: Int -> Complex m a
+  Lots            :: a -> b -> c -> d -> e -> f -> Complex m ()
+  Nested          :: Maybe b -> Complex m (Maybe a)
+  MultiNested     :: (Maybe a, [b]) -> Complex m (Maybe a, [b])
+  Existential     :: (forall e. e -> Maybe e) -> Complex m a
+  LotsNested      :: Maybe a -> [b] -> (c, c) -> Complex m (a, b, c)
+  Dict            :: Ord a => a -> Complex m a
+  MultiDict       :: (Eq a, Ord b, Enum a, Num c)
+                  => a -> b -> c -> Complex m ()
+  IndexedMono     :: f 0 -> Complex m Int
+  IndexedPoly     :: forall f (n :: Nat) m . f n -> Complex m (f (n + 1))
+  IndexedPolyDict :: KnownNat n => f n -> Complex m Int
+
+makeSem ''Complex
+
+data HOEff m a where
+  EffArgMono :: m () -> HOEff m ()
+  EffArgPoly :: m a -> HOEff m a
+  EffArgComb :: m a -> (m a -> m b) -> HOEff m b
+  EffRank2   :: (forall x. m x -> m (Maybe x)) -> HOEff m a
+
+makeSem ''HOEff
+
+data ComplexEffArgs b c m a where
+  EffMono     :: Int -> ComplexEffArgs Int String m Bool
+  EffPoly1    :: a -> ComplexEffArgs a b m a
+  EffPoly2    :: a -> ComplexEffArgs a (Maybe a) m Bool
+  EffPolyFree :: String -> ComplexEffArgs a b m Int
+  EffSame1    :: ComplexEffArgs a a m a
+  EffSame2    :: ComplexEffArgs b b m a
+  EffHO       :: m b -> ComplexEffArgs b Int m String
+
+makeSem ''ComplexEffArgs
+
+data HKEffArgs f g m a where
+  HKRank2 :: (forall x . f x -> g x) -> HKEffArgs f g m a
+
+makeSem ''HKEffArgs
+
+-- 'makeSem' input names -----------------------------------------------------
+
+data ByCon m a where
+  ByConC :: Int -> ByCon m String
+
+makeSem 'ByConC
+
+data ByField m a where
+  ByFieldC :: { byFieldCF :: Int } -> ByField m Int
+
+makeSem 'byFieldCF