-- SPDX-FileCopyrightText: 2020 Tocqueville Group -- -- SPDX-License-Identifier: LicenseRef-MIT-TQ -- -- | Generic deriving with unbalanced trees. module Util.CustomGeneric ( -- * Custom Generic strategies GenericStrategy , withDepths , rightBalanced , leftBalanced , rightComb , leftComb , haskellBalanced -- ** Entries reordering , reorderingConstrs , reorderingFields , reorderingData , alphabetically , leaveUnnamedFields , forbidUnnamedFields -- * Depth usage helpers , cstr , fld -- * Instance derivation , customGeneric -- * Helpers , fromDepthsStrategy -- * Internals , reifyDataType , deriveFullType , customGeneric' ) where import Control.Lens (traversed) import Generics.Deriving.TH (makeRep0Inline) import qualified GHC.Generics as G import Language.Haskell.TH import Util.Generic (mkGenericTree) import Util.TH (lookupTypeNameOrFail) ---------------------------------------------------------------------------- -- Simple type synonyms ---------------------------------------------------------------------------- -- | Simple tuple specifying the depth of a constuctor and a list of depths -- for its fields. -- -- This is used as a way to specify the tree topology of the Generic instance -- to derive. type CstrDepth = (Natural, [Natural]) -- | Simple tuple that defines the "shape" of a constructor: it's name and number -- of fields. Used only in this module. type CstrShape = (Name, Int) -- | Simple tuple that carries basic info about a constructor: it's name, -- number of its fields and their names. Used only in this module. type CstrNames = (Name, Int, Maybe [Name]) -- | Type of a strategy to derive 'G.Generic' instances. data GenericStrategy = GenericStrategy { gsEvalDepths :: [CstrShape] -> Q [CstrDepth] -- ^ Given the 'CstrShape's for given datatype, -- return the 'CstrDepth's for it. -- This function should when possible make checks and 'fail', using the -- constructors' 'Name' provided by the 'CstrShape'. , gsReorderCstrsOn :: forall a. [(Text, a)] -> Q [a] -- ^ Reorder constructors given their names. , gsReorderFieldsOn :: forall a. Either [a] [(Text, a)] -> Q [a] -- ^ Reorder fields given their names, the argument depends on whether -- fields are part of record (and thus named) or not (unnamed). } -- | Defines how to reorder fields or constructors. type EntriesReorder = forall a. [(Text, a)] -> Q [a] -- | Defines how to reorder fields when their names are unknown. type UnnamedEntriesReorder = forall a. [a] -> Q [a] -- | Simple type synonym used (internally) between functions, basically extending -- 'CstrDepth' with the 'Name's of the constructor and its fields. -- For fields it carries both names in the original order and in the order specified -- by the strategy (and the latter is paired with depths). data NamedCstrDepths = NCD { ncdCstrDepth :: Natural -- ^ Constructor's depth , ncdCstrName :: Name -- ^ Constructor's name , ncdOrigFieldNames :: [Name] -- ^ Names of constructor fields in the original order. , ncdFields :: [(Natural, Name)] -- ^ Names and depths of constructor fields after the reordering. } -- | Reorders entries corresponding to constructors (@a@) and fields (@b@) -- according to some rule. type EntriesTransp = forall a b. [a] -> Q [([b] -> Q [b], a)] ---------------------------------------------------------------------------- -- Generic strategies ---------------------------------------------------------------------------- -- | In this strategy the desired depths of contructors (in the type tree) and -- fields (in each constructor's tree) are provided manually and simply checked -- against the number of actual constructors and fields. withDepths :: [CstrDepth] -> GenericStrategy withDepths treeDepths = simpleGenericStrategy $ \cstrShape -> do when (length treeDepths /= length cstrShape) $ fail "Number of contructors' depths does not match number of data contructors." forM_ (zip (map snd treeDepths) cstrShape) $ \(fDepths, (constrName, fldNum)) -> when (length fDepths /= fldNum) . fail $ "Number of fields' depths does not match number of field for data " <> "constructor: " <> show constrName return treeDepths -- | Strategy to make right-balanced instances (both in constructors and fields). -- -- This will try its best to produce a flat tree: -- -- * the balances of all leaves differ no more than by 1; -- * leaves at left will have equal or lesser depth than leaves at right. rightBalanced :: GenericStrategy rightBalanced = fromDepthsStrategy makeRightBalDepths -- | Strategy to make left-balanced instances (both in constructors and fields). -- -- This is the same as symmetrically mapped 'rightBalanced'. leftBalanced :: GenericStrategy leftBalanced = fromDepthsStrategy (reverse . makeRightBalDepths) -- | Strategy to make fully right-leaning instances (both in constructors and fields). rightComb :: GenericStrategy rightComb = fromDepthsStrategy (reverse . makeLeftCombDepths) -- | Strategy to make fully left-leaning instances (both in constructors and fields). leftComb :: GenericStrategy leftComb = fromDepthsStrategy makeLeftCombDepths -- | Strategy to make Haskell's Generics-like instances -- (both in constructors and fields). -- -- This is similar to 'rightBalanced', except for the "flat" part: -- -- * for each node, size of the left subtree is equal or less by one than -- size of the right subtree. -- -- This strategy matches A1.1. -- -- @customGeneric "T" haskellBalanced@ is equivalent to mere -- @deriving stock Generic T@. haskellBalanced :: GenericStrategy haskellBalanced = fromDepthsStrategy makeHaskellDepths -- Order modifiers ---------------------------------------------------------------------------- -- | Modify given strategy to reorder constructors. -- -- The reordering will take place before depths are evaluated and structure -- of generic representation is formed. -- -- Example: @reorderingConstrs alphabetically rightBalanced@. reorderingConstrs :: EntriesReorder -> GenericStrategy -> GenericStrategy reorderingConstrs reorder gs = gs { gsReorderCstrsOn = reorder } -- | Modify given strategy to reorder fields. -- -- Same notes as for 'reorderingConstrs' apply here. -- -- Example: @reorderingFields forbidUnnamedFields alphabetically rightBalanced@. reorderingFields :: UnnamedEntriesReorder -> EntriesReorder -> GenericStrategy -> GenericStrategy reorderingFields reorderUnnamed reorder gs = gs { gsReorderFieldsOn = either reorderUnnamed reorder } -- | Modify given strategy to reorder constructors and fields. -- -- Same notes as for 'reorderingConstrs' apply here. -- -- Example: @reorderingData forbidUnnamedFields alphabetically rightBalanced@. reorderingData :: UnnamedEntriesReorder -> EntriesReorder -> GenericStrategy -> GenericStrategy reorderingData reorderUnnamed reorder = reorderingFields reorderUnnamed reorder . reorderingConstrs reorder -- | Sort entries by name alphabetically. alphabetically :: EntriesReorder alphabetically = pure . map snd . sortWith fst -- | Leave unnamed fields intact, without any reordering. leaveUnnamedFields :: UnnamedEntriesReorder leaveUnnamedFields = pure -- | Fail in case records are unnamed and we cannot figure -- out the necessary reordering. forbidUnnamedFields :: UnnamedEntriesReorder forbidUnnamedFields fields = if length fields <= 1 then return fields else fail "Encountered unnamed fields, cannot apply reordering" ---------------------------------------------------------------------------- -- Generic strategies' builders ---------------------------------------------------------------------------- -- | Construct a strategy that only constructs Generic instance of given -- form, without e.g. changing the order of entries. simpleGenericStrategy :: ([CstrShape] -> Q [CstrDepth]) -> GenericStrategy simpleGenericStrategy mkDepths = GenericStrategy { gsEvalDepths = mkDepths , gsReorderCstrsOn = pure . map snd , gsReorderFieldsOn = pure . either id (map snd) } -- | Helper to make a strategy that created depths for constructor and fields -- in the same way, just from their number. -- -- The provided function @f@ must satisfy the following rules: -- -- * @length (f n) ≡ n@ -- * @sum $ (\x -> 2 ^^ (-x)) <$> f n ≡ 1@ (unless @n = 0@) fromDepthsStrategy :: (Int -> [Natural]) -> GenericStrategy fromDepthsStrategy dStrategy = simpleGenericStrategy $ \cShapes -> return $ zip (dStrategy $ length cShapes) $ map (dStrategy . view _2) cShapes makeRightBalDepths :: Int -> [Natural] makeRightBalDepths n = foldr (const addRightBalDepth) [] [1..n] where addRightBalDepth :: [Natural] -> [Natural] addRightBalDepth = \case [] -> [0] [x] -> [x + 1, x + 1] (x : y : xs) | x == y -> x : addRightBalDepth (x : xs) (_ : y : xs) -> y : y : y : xs makeLeftCombDepths :: Int -> [Natural] makeLeftCombDepths 0 = [] makeLeftCombDepths n = map fromIntegral $ (n - 1) : [n - 1, n - 2..1] makeHaskellDepths :: Int -> [Natural] makeHaskellDepths n = case nonEmpty (replicate n [0]) of Nothing -> [] Just leaves -> mkGenericTree (\_ l r -> map succ (l ++ r)) leaves ---------------------------------------------------------------------------- -- Depth usage helpers ---------------------------------------------------------------------------- -- | Helper for making a constructor depth. -- -- Note that this is only intended to be more readable than directly using a -- tuple with 'withDepths' and for the ability to be used in places where -- @RebindableSyntax@ overrides the number literal resolution. cstr :: forall n. KnownNat n => [Natural] -> CstrDepth cstr flds = (natVal (Proxy @n), flds) -- | Helper for making a field depth. -- -- Note that this is only intended to be more readable than directly using a -- tuple with 'withDepths' and for the ability to be used in places where -- @RebindableSyntax@ overrides the number literal resolution. fld :: forall n. KnownNat n => Natural fld = natVal $ Proxy @n ---------------------------------------------------------------------------- -- Instance derivation ---------------------------------------------------------------------------- {-# ANN module ("HLint: ignore Use snd" :: Text) #-} -- | Derives the 'G.Generic' instance for a type given its name and a -- 'GenericStrategy' to use. -- -- The strategy is used to calculate the depths of the data-type constructors -- and each constructors' fields. -- -- The depths are used to generate the tree of the 'G.Generic' representation, -- allowing for a custom one, in contrast with the one derived automatically. -- -- This only supports "plain" @data@ types (no GADTs, no @newtype@s, etc.) and -- requires the depths to describe a fully and well-defined tree (see 'unbalancedFold'). -- -- For example, this is valid (and uses the 'withDepths' strategy with the 'cstr' -- and 'fld' helpers) and results in a balanced instance, equivalent to the -- auto-derived one: -- -- @@@ -- data CustomType a -- = CustomUp Integer Integer -- | CustomMid {unMid :: Natural} -- | CustomDown a -- | CustomNone -- -- $(customGeneric "CustomType" $ withDepths -- [ cstr @2 [fld @1, fld @1] -- , cstr @2 [fld @0] -- , cstr @2 [fld @0] -- , cstr @2 [] -- ]) -- @@@ -- -- and this is a valid, but fully left-leaning one: -- -- @@@ -- $(customGeneric "CustomType" $ withDepths -- [ cstr @3 [fld @1, fld @1] -- , cstr @3 [fld @0] -- , cstr @2 [fld @0] -- , cstr @1 [] -- ]) -- @@@ -- -- and, just as a demonstration, this is the same fully left-leaning one, but -- made using the simpler 'leftComb' strategy: -- -- @@@ -- $(customGeneric "CustomType" leftComb) -- @@@ -- -- Developers are welcome to provide their own derivation strategies, -- and some useful strategies can be found outside of this module by -- 'GenericStrategy' signature. customGeneric :: String -> GenericStrategy -> Q [Dec] customGeneric typeStr genStrategy = do -- Implementor's note: -- -- Instead of using a name literal (@customGeneric ''T@), we use a string (@customGeneric "T"@) -- and then 'lookupTypeName' for the following reasons: -- -- 1. We can control the error message when 'lookupTypeName' doesn't find the type in scope (as opposed to @''T@) -- 2. Most importantly, this was made with Indigo in mind, where we try as much as -- possible to use a simple syntax (to appeal to a broader audience) and so to avoid -- using more obscure Haskell syntax (like @''T@). -- reify the data type (typeName, _, mKind, vars, constructors) <- lookupTypeNameOrFail typeStr >>= reifyDataType -- obtain info about its constructor and desired tree derivedType <- deriveFullType typeName mKind vars customGeneric' Nothing typeName derivedType constructors genStrategy -- | If a 'Rep' type is given, this function will generate a new 'Generic' instance with it, -- and generate the appropriate "to" and "from" methods. -- -- Otherwise, it'll generate a new 'Rep' instance as well. customGeneric' :: Maybe Type -> Name -> Type -> [Con] -> GenericStrategy -> Q [Dec] customGeneric' maybeRepType typeName derivedType constructors genStrategy = do cNames <- cstrNames constructors let cReordering :: EntriesTransp cReordering = reorderCstrs genStrategy cNames let cShapes = cNames <&> \(name, fNum, _) -> (name, fNum) cShapesSorted <- cReordering cShapes <&> map \(_fReorder, cShape) -> cShape treeDepths <- gsEvalDepths genStrategy cShapesSorted weightedConstrs <- makeWeightedConstrs cReordering treeDepths cShapesSorted -- If no 'Rep' type was given, derive one. let repType = maybe (makeUnbalancedRep typeName treeDepths cReordering (pure derivedType)) pure maybeRepType -- produce the Generic instance res <- instanceD (pure []) (conT ''G.Generic `appT` pure derivedType) [ tySynInstD . tySynEqn Nothing (conT ''G.Rep `appT` pure derivedType) $ repType , makeUnbalancedFrom weightedConstrs , makeUnbalancedTo weightedConstrs ] return [res] -- | Apply a reordering strategy. -- -- This uses given @[CstrNames]@ to understand how constructors and their -- fields should be reordered, and applies the same transposition to entries -- within 'EntriesTransp'. reorderCstrs :: GenericStrategy -> [CstrNames] -> EntriesTransp reorderCstrs GenericStrategy{..} cNames = \cstrEntries -> gsReorderCstrsOn $ zip cNames cstrEntries <&> \(cstrName@(name, _, _), cstrEntry) -> (origName name, (fieldsReorder cstrName, cstrEntry)) where fieldsReorder :: CstrNames -> [b] -> Q [b] fieldsReorder (_, _, mFieldNames) = \fieldEntries -> do gsReorderFieldsOn $ maybe Left (Right ... zip . map origName) mFieldNames fieldEntries -- | Reifies info from a type name (given as a 'String'). -- The lookup happens from the current splice's scope (see 'lookupTypeName') and -- the only accepted result is a "plain" data type (no GADTs). reifyDataType :: Name -> Q (Name, Cxt, Maybe Kind, [TyVarBndr], [Con]) reifyDataType typeName = do typeInfo <- reify typeName case typeInfo of TyConI (DataD decCxt typeName' vars mKind constrs _) -> return (typeName', decCxt, mKind, vars, constrs) _ -> fail $ "Only plain datatypes are supported for derivation, but '" <> show typeName <> "' instead reifies to:\n" <> show typeInfo -- | Derives, as well as possible, a type definition from its name, its kind -- (where known) and its variables. deriveFullType :: Name -> Maybe Kind -> [TyVarBndr] -> TypeQ deriveFullType tName mKind = addTypeSig . foldl appT (conT tName) . makeVarsType where addTypeSig :: TypeQ -> TypeQ addTypeSig = flip sigT $ fromMaybe StarT mKind makeVarsType :: [TyVarBndr] -> [TypeQ] makeVarsType = map $ \case PlainTV vName -> varT vName KindedTV vName kind -> sigT (varT vName) kind -- | Extract the info for each of the given constructors. cstrNames :: [Con] -> Q [CstrNames] cstrNames constructors = forM constructors $ \case NormalC name lst -> return (name, length lst, Nothing) RecC name lst -> return (name, length lst, Just $ lst ^.. traversed . _1) InfixC _ name _ -> return (name, 2, Nothing) constr -> fail $ "Unsupported constructor: " <> show constr -- | Combines depths with constructors, 'fail'ing in case of mismatches, and -- generates 'Name's for the constructors' arguments. makeWeightedConstrs :: EntriesTransp -> [CstrDepth] -> [CstrShape] -> Q [NamedCstrDepths] makeWeightedConstrs cReorder treeDepths cShapes = do reorderedShapes <- cReorder cShapes forM (zip treeDepths reorderedShapes) $ \((cDepth, fDepths), (fReorder, (cName, fNum))) -> do fieldVarsNames <- forM [0 .. fNum - 1] \i -> newName ("v" <> show i) reorderedFieldVarNames <- fReorder fieldVarsNames return NCD { ncdCstrDepth = cDepth , ncdCstrName = cName , ncdOrigFieldNames = fieldVarsNames , ncdFields = zip fDepths reorderedFieldVarNames } -- | Creates the 'G.Rep' type for an unbalanced 'G.Generic' instance, for a type -- given its name, constructors' depths and derived full type. -- -- Note: given that these types definition can be very complex to generate, -- especially in the metadata, here we let @generic-deriving@ make a balanced -- value first (see 'makeRep0Inline') and then de-balance the result. makeUnbalancedRep :: Name -> [CstrDepth] -> EntriesTransp -> TypeQ -> TypeQ makeUnbalancedRep typeName treeDepths reorderConstrs derivedType = do -- let generic-deriving create the balanced type first balRep <- makeRep0Inline typeName derivedType -- separate the top-most type metadata from the constructors' trees (typeMd, constrTypes) <- dismantleGenericTree [t| G.C1 |] balRep -- for each of the constructor's trees reorderedConstrTypes <- reorderConstrs constrTypes unbalConstrs <- forM (zip reorderedConstrTypes treeDepths) $ \((reorderFields, constrType), treeDepth) -> case treeDepth of (n, []) -> -- when there are no fields there is no tree to unbalance return (n, constrType) (n, fieldDepths) -> do -- separate the top-most constructor metadata from the fields' trees (constrMd, fieldTypes) <- dismantleGenericTree [t| G.S1 |] constrType -- build the unbalanced tree of fields reorderedFieldTypes <- reorderFields fieldTypes unbalConstRes <- unbalancedFold (zip fieldDepths reorderedFieldTypes) (appT . appT (conT ''(G.:*:))) -- return the new unbalanced constructor return (n, AppT constrMd unbalConstRes) -- build the unbalanced tree of constructors and rebuild the type appT (pure typeMd) $ unbalancedFold unbalConstrs (appT . appT (conT ''(G.:+:))) -- | Breaks down a tree of @Generic@ types from the contructor of "nodes" and -- the constructor for "leaves" metadata. -- -- This expects (and should always be the case) the "root" to be a @Generic@ -- metadata contructor, which is returned in the result alongside the list of -- leaves (in order). dismantleGenericTree :: TypeQ -> Type -> Q (Type, [Type]) dismantleGenericTree leafMetaQ (AppT meta nodes) = do leafMeta <- leafMetaQ let collectLeafsTypes :: Type -> [Type] collectLeafsTypes tp = case tp of f `AppT` _ `AppT` _ | f == leafMeta -> [tp] AppT a b -> collectLeafsTypes a <> collectLeafsTypes b _ -> [] return (meta, collectLeafsTypes nodes) dismantleGenericTree _ x = fail $ "Unexpected lack of Generic Metadata: " <> show x -- | Create the unbalanced 'G.from' fuction declaration for a type starting from -- its list of weighted constructors. makeUnbalancedFrom :: [NamedCstrDepths] -> DecQ makeUnbalancedFrom wConstrs = do (cPatts, cDepthExp) <- fmap unzip . forM wConstrs $ \(NCD cDepth cName wOrigFields wFields) -> do fDepthExp <- forM wFields $ \(fDepth, fName) -> do -- make expression to asseble a Generic Field from its variable fExpr <- appE [| G.M1 |] . appE [| G.K1 |] $ varE fName return (fDepth, fExpr) -- make pattern for this constructor fPatts <- mapM varP wOrigFields let cPatt = ConP cName fPatts -- make expression to assemble its fields as an isolated Generic Constructor cExp <- appE [| G.M1 |] $ case fDepthExp of [] -> conE 'G.U1 _ -> unbalancedFold fDepthExp (appE . appE [| (G.:*:) |]) return (cPatt, (cDepth, [cExp])) -- make expressions to assemble all Generic Constructors cExps <- mapQ (appE [| G.M1 |]) $ unbalancedFold cDepthExp $ \xs ys -> (<>) <$> mapQ (appE [| G.L1 |]) xs <*> mapQ (appE [| G.R1 |]) ys -- make function definition funD 'G.from $ zipWith (\p e -> clause [pure p] (normalB $ pure e) []) cPatts cExps -- | Create the unbalanced 'G.to' fuction declaration for a type starting from -- its list of weighted constructors. makeUnbalancedTo :: [NamedCstrDepths] -> DecQ makeUnbalancedTo wConstrs = do (cExps, cDepthPat) <- fmap unzip . forM wConstrs $ \(NCD cDepth cName wOrigFields wFields) -> do fDepthPat <- forM wFields $ \(fDepth, fName) -> do -- make pattern for a Generic Field from its variable fPatt <- conP1 'G.M1 . conP1 'G.K1 $ varP fName return (fDepth, fPatt) -- make pattern for this isolated Generic Constructor cPatt <- conP1 'G.M1 $ case fDepthPat of [] -> conP 'G.U1 [] _ -> unbalancedFold fDepthPat (conP2 '(G.:*:)) -- make expression to assemble this constructor fExps <- mapM varE wOrigFields let cExp = foldl AppE (ConE cName) fExps return (cExp, (cDepth, [cPatt])) -- make patterns for all Generic Constructors cPatts <- mapQ (conP1 'G.M1) $ unbalancedFold cDepthPat $ \xs ys -> (<>) <$> mapQ (conP1 'G.L1) xs <*> mapQ (conP1 'G.R1) ys -- make function definition funD 'G.to $ zipWith (\p e -> clause [pure p] (normalB $ pure e) []) cPatts cExps -- | Recursively aggregates the values in the given list by merging (with the -- given function) the ones that are adjacent and with the same depth. -- -- This will fail for every case in which the list cannot be folded into a single -- 0-depth value. unbalancedFold :: forall a. Eq a => [(Natural, a)] -> (Q a -> Q a -> Q a) -> Q a unbalancedFold lst f = unbalancedFoldRec lst >>= \case [(0, result)] -> return result [(n, _)] -> fail $ "Resulting unbalanced tree has a single root, but of depth " <> show n <> " instead of 0. Check your depths definitions." _ -> fail $ "Cannot create a tree from nodes of depths: " <> show (map fst lst) <> ". Check your depths definitions." where unbalancedFoldRec :: [(Natural, a)] -> Q [(Natural, a)] unbalancedFoldRec xs = do ys <- unbalancedFoldSingle xs if xs == ys then return xs else unbalancedFoldRec ys unbalancedFoldSingle :: [(Natural, a)] -> Q [(Natural, a)] unbalancedFoldSingle = \case [] -> return [] (dx, x) : (dy, y) : xs | dx == dy -> do dxy <- f (pure x) (pure y) return $ (dx - 1, dxy) : xs x : xs -> do ys <- unbalancedFoldSingle xs return (x : ys) ---------------------------------------------------------------------------- -- Utility functions ---------------------------------------------------------------------------- conP1 :: Name -> PatQ -> PatQ conP1 name pat = conP name [pat] conP2 :: Name -> PatQ -> PatQ -> PatQ conP2 name pat1 pat2 = conP name [pat1, pat2] mapQ :: (Q a -> Q a) -> Q [a] -> Q [a] mapQ f qlst = qlst >>= mapM (f . pure) -- | Original name of a constructor or field. origName :: Name -> Text origName = toText . nameBase