Compute a mapped type family's image at the use site

aeson-typescript.cabal

@@ -1,6 +1,6 @@
 cabal-version: 1.12
 
--- This file has been generated from package.yaml by hpack version 0.38.0.
+-- This file has been generated from package.yaml by hpack version 0.39.1.
 --
 -- see: https://github.com/sol/hpack
 
@@ -104,6 +104,7 @@ test-suite aeson-typescript-tests
       UntaggedNoTagSingleConstructors
       UntaggedTagSingleConstructors
       UnwrapUnaryRecords
+      UseSiteFamilyImage
       Util
       Util.Aeson
       Data.Aeson.TypeScript.Formatting

src/Data/Aeson/TypeScript/Lookup.hs

@@ -1,10 +1,16 @@
 {-# LANGUAGE CPP #-}
-{-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
 
-module Data.Aeson.TypeScript.Lookup where
+module Data.Aeson.TypeScript.Lookup (
+  deriveTypeScriptLookupType
+
+  , getClosedTypeFamilyImage
+  , getClosedTypeFamilyInterfaceDecl'
+  )
+  where
 
 import Control.Monad
 import Data.Aeson.TypeScript.Instances ()
@@ -14,6 +20,7 @@ import qualified Data.List as L
 import Data.Proxy
 import Data.String.Interpolate
 import Language.Haskell.TH hiding (stringE)
+import Language.Haskell.TH.Datatype.TyVarBndr ()
 import qualified Language.Haskell.TH.Lib as TH
 
 
@@ -33,7 +40,10 @@ deriveTypeScriptLookupType name declNameStr = do
     _ -> fail [i|Expected a closed type family; got #{info}|]
 
 getClosedTypeFamilyInterfaceDecl :: Name -> [TySynEqn] -> Q Exp
-getClosedTypeFamilyInterfaceDecl name eqns = do
+getClosedTypeFamilyInterfaceDecl name eqns = getClosedTypeFamilyInterfaceDecl' (nameBase name) eqns
+
+getClosedTypeFamilyInterfaceDecl' :: String -> [TySynEqn] -> Q Exp
+getClosedTypeFamilyInterfaceDecl' name eqns = do
   fields <- forM eqns $ \case
 #if MIN_VERSION_template_haskell(2,15,0)
     TySynEqn Nothing (AppT (ConT _) (ConT arg)) result -> do
@@ -46,7 +56,7 @@ getClosedTypeFamilyInterfaceDecl name eqns = do
 #endif
     x -> fail [i|aeson-typescript doesn't know yet how to handle this type family equation when generating interface declaration: '#{x}'|]
 
-  [| TSInterfaceDeclaration $(TH.stringE $ nameBase name) [] (L.sortBy (compare `on` fieldName) $(listE $ fmap return fields)) Nothing |]
+  [| TSInterfaceDeclaration $(TH.stringE name) [] (L.sortBy (compare `on` fieldName) $(listE $ fmap return fields)) Nothing |]
 
 getClosedTypeFamilyImage :: [TySynEqn] -> Q [Type]
 getClosedTypeFamilyImage eqns = do

src/Data/Aeson/TypeScript/TH.hs

@@ -145,7 +145,7 @@
   , T9(..)
   , T10(..)
 
   , module Data.Aeson.TypeScript.Instances
   ) where
 
 import Control.Monad
@@ -183,9 +183,18 @@
   datatypeInfo' <- reifyDatatype name
   assertExtensionsTurnedOn datatypeInfo'
 
-  -- Figure out what the generic variables are
+  -- Figure out what the generic variables are. A variable of kind * is one; so is one that
+  -- keys a mapped type family in some field, whatever its kind, since the family application
+  -- becomes an indexed access on it.
+  let fieldTypes = mconcat [constructorFields ci | ci <- datatypeCons datatypeInfo']
+  let keysAMappedFamily n = or [True | typ <- fieldTypes
+                                     , (ConT famName, VarT keyVar : _) <- [splitAppsT typ]
+                                     , keyVar == n
+                                     , famName `L.elem` typeFamiliesToMapToTypeScript extraOptions]
   let eligibleGenericVars = catMaybes $ flip fmap (getDataTypeVars datatypeInfo') $ \case
         SigT (VarT n) StarT -> Just n
+        SigT (VarT n) _ | keysAMappedFamily n -> Just n
+        VarT n | keysAMappedFamily n -> Just n
         _ -> Nothing
   let varsAndTVars = case eligibleGenericVars of
         [] -> []
@@ -291,7 +300,7 @@
 
 #if MIN_VERSION_aeson(0,10,0)
       | unwrapUnaryRecords options && (isSingleRecordConstructor ci) -> do
           let [typ] = constructorFields ci
           stringExp <- lift $ [|getTypeScriptTypeOrOptionalNull (Proxy :: Proxy $(return typ))|]
           alternatives <- lift [|TSTypeAlternatives $(TH.stringE interfaceName)
                                                     $(genericVariablesListExpr True genericVariables)

src/Data/Aeson/TypeScript/Transform.hs

@@ -1,48 +1,76 @@
 {-# LANGUAGE CPP #-}
-{-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
 
 
 module Data.Aeson.TypeScript.Transform (
   transformTypeFamilies
   ) where
 
+import Control.Monad (unless)
 import Control.Monad.Writer
 import Data.Aeson.TypeScript.Lookup
 import Data.Aeson.TypeScript.Types
+import Data.Aeson.TypeScript.Util (splitAppsT, typeFamilyTableName)
 import qualified Data.List as L
+import qualified Data.Map as M
+import Data.Maybe (fromMaybe, isJust)
+import qualified Data.Set as S
 import Data.Typeable
 import Language.Haskell.TH hiding (stringE)
+import Language.Haskell.TH.Datatype (applySubstitution)
 import qualified Language.Haskell.TH.Lib as TH
+import Language.Haskell.TH.Syntax (getQ, putQ)
 
 #if !MIN_VERSION_base(4,11,0)
 import Data.Monoid
 #endif
 
 
+-- | Tables emitted so far in this module, so that repeated use sites of one family agree on
+-- a single declaration. 'Q' state persists across splices within a module.
+newtype EmittedFamilyTables = EmittedFamilyTables (S.Set String)
+
 -- | Search the given type for type families. For each one found, emit a declaration for a new
 -- corresponding concrete type and a TypeScript instance for it which emits a lookup type.
 -- Then, replace all occurrences of the given type family with the concrete type in the return value.
 -- Thus the type becomes "de-family-ified".
 transformTypeFamilies :: ExtraTypeScriptOptions -> Type -> WriterT [ExtraDeclOrGenericInfo] Q Type
-transformTypeFamilies eo@(ExtraTypeScriptOptions {..}) (AppT (ConT name) typ)
-  | name `L.elem` typeFamiliesToMapToTypeScript = lift (reify name) >>= \case
-      FamilyI (ClosedTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _) eqns) _ -> handle typeFamilyName eqns
+transformTypeFamilies eo@(ExtraTypeScriptOptions {..}) wholeType
+  | (ConT name, keyArg : trailing) <- splitAppsT wholeType
+  , name `L.elem` typeFamiliesToMapToTypeScript = lift (reify name) >>= \case
+      FamilyI (ClosedTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _) eqns) _ -> handle name keyArg trailing typeFamilyName eqns
 
 #if MIN_VERSION_template_haskell(2,15,0)
-      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) decs -> handle typeFamilyName [eqn | TySynInstD eqn <- decs]
+      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) decs -> handle name keyArg trailing typeFamilyName [eqn | TySynInstD eqn <- decs]
 #else
-      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) decs -> handle typeFamilyName [eqn | TySynInstD _name eqn <- decs]
+      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) decs -> handle name keyArg trailing typeFamilyName [eqn | TySynInstD _name eqn <- decs]
 #endif
 
-      _ -> AppT (ConT name) <$> transformTypeFamilies eo typ
-  | otherwise = AppT (ConT name) <$> transformTypeFamilies eo typ
-        where
-          handle :: Name -> [TySynEqn] -> WriterT [ExtraDeclOrGenericInfo] Q Type
-          handle typeFamilyName eqns = do
-            name' <- lift $ newName (nameBase typeFamilyName <> "'")
+      _ -> descend wholeType
+  where
+          descend (AppT typ1 typ2) = AppT <$> transformTypeFamilies eo typ1 <*> transformTypeFamilies eo typ2
+          descend (SigT typ kind) = flip SigT kind <$> transformTypeFamilies eo typ
+          descend (ParensT typ) = ParensT <$> transformTypeFamilies eo typ
+          descend typ = return typ
+
+          handle :: Name -> Type -> [Type] -> Name -> [TySynEqn] -> WriterT [ExtraDeclOrGenericInfo] Q Type
+          handle _name keyArg trailing typeFamilyName eqnsRaw = do
+            eqns <- lift $ mapM (substituteTrailingArgs trailing) eqnsRaw
+            let tableName = typeFamilyTableName typeFamilyName trailing
+            -- Deterministic and named after the table, so that repeated use sites of one
+            -- family agree on it rather than each declaring their own.
+            let name' = mkName (tableName <> "'")
+            alreadyDeclared <- lift $ do
+              inScope <- isJust <$> lookupTypeName (tableName <> "'")
+              EmittedFamilyTables emitted <- fromMaybe (EmittedFamilyTables mempty) <$> getQ
+              case inScope || S.member tableName emitted of
+                True -> return True
+                False -> do
+                  putQ (EmittedFamilyTables (S.insert tableName emitted))
+                  return False
 
             f <- lift $ newName "f"
 #if MIN_VERSION_template_haskell(2,21,0)
@@ -52,19 +80,19 @@
 #else
             let inst1 = DataD [] name' [PlainTV f] Nothing [] []
 #endif
-            tell [ExtraTopLevelDecs [inst1]]
+            unless alreadyDeclared $ tell [ExtraTopLevelDecs [inst1]]
 
             imageTypes <- lift $ getClosedTypeFamilyImage eqns
             inst2 <- lift $ [d|instance (Typeable g, TypeScript g) => TypeScript ($(conT name') g) where
-                                 getTypeScriptType _ = $(TH.stringE $ nameBase name) <> "[" <> (getTypeScriptType (Proxy :: Proxy g)) <> "]"
-                                 getTypeScriptDeclarations _ = [$(getClosedTypeFamilyInterfaceDecl name eqns)]
+                                 getTypeScriptType _ = $(TH.stringE tableName) <> "[" <> (getTypeScriptType (Proxy :: Proxy g)) <> "]"
+                                 getTypeScriptDeclarations _ = [$(getClosedTypeFamilyInterfaceDecl' tableName eqns)]
                                  getParentTypes _ = $(listE [ [|TSType (Proxy :: Proxy $(return x))|] | x <- imageTypes])
                             |]
-            tell [ExtraTopLevelDecs inst2]
+            unless alreadyDeclared $ tell [ExtraTopLevelDecs inst2]
 
             tell [ExtraParentType (AppT (ConT name') (ConT ''T))]
 
-            ret <- transformTypeFamilies eo (AppT (ConT name') typ)
+            ret <- transformTypeFamilies eo (AppT (ConT name') keyArg)
             tell [ExtraConstraint (AppT (ConT ''TypeScript) ret)]
             return ret
 transformTypeFamilies eo (AppT typ1 typ2) = AppT <$> transformTypeFamilies eo typ1 <*> transformTypeFamilies eo typ2
@@ -77,3 +105,29 @@
 transformTypeFamilies eo (ImplicitParamT s typ) = ImplicitParamT s <$> transformTypeFamilies eo typ
 #endif
 transformTypeFamilies _ typ = return typ
+
+-- | Fix a family's trailing arguments to the types a use site applies it to, so that an
+-- equation whose result is one of those arguments gets a concrete result.
+--
+-- @Chosen f x y@, with @Chosen Terse x y = x@, has no image on its own. Applied to
+-- @Chosen f Int String@ it becomes @Chosen Terse Int String = Int@, which does.
+substituteTrailingArgs :: [Type] -> TySynEqn -> Q TySynEqn
+#if MIN_VERSION_template_haskell(2,15,0)
+substituteTrailingArgs [] eqn = return eqn
+substituteTrailingArgs args (TySynEqn bndrs lhs result) = case splitLhs lhs of
+  Just (famName, key : rest) | length rest == length args -> do
+    let substitution = M.fromList [(n, t) | (VarT n, t) <- zip rest args]
+    return $ TySynEqn Nothing (AppT (ConT famName) key) (applySubstitution substitution result)
+  _ -> return (TySynEqn bndrs lhs result)
+  where
+    splitLhs :: Type -> Maybe (Name, [Type])
+    splitLhs (ConT n) = Just (n, [])
+    splitLhs (AppT f x) = do
+      (n, xs) <- splitLhs f
+      return (n, xs <> [x])
+    splitLhs (SigT t _) = splitLhs t
+    splitLhs (ParensT t) = splitLhs t
+    splitLhs _ = Nothing
+#else
+substituteTrailingArgs _ eqn = return eqn
+#endif

src/Data/Aeson/TypeScript/TypeManipulation.hs

@@ -12,6 +12,7 @@
 
 import Control.Monad.Writer
 import Data.Aeson.TypeScript.Types
+import Data.Aeson.TypeScript.Util (splitAppsT, typeFamilyTableName)
 import qualified Data.List as L
 import Language.Haskell.TH
 
@@ -21,16 +22,15 @@
 
 
 searchForConstraints :: ExtraTypeScriptOptions -> Type -> Name -> WriterT [GenericInfo] Q ()
-searchForConstraints eo@(ExtraTypeScriptOptions {..}) (AppT (ConT name) typ) var
-  | typ == VarT var && (name `L.elem` typeFamiliesToMapToTypeScript) = lift (reify name) >>= \case
-      FamilyI (ClosedTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _) _) _ -> do
-        tell [GenericInfo var (TypeFamilyKey typeFamilyName)]
-        searchForConstraints eo typ var
-      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) _ -> do
-        tell [GenericInfo var (TypeFamilyKey typeFamilyName)]
-        searchForConstraints eo typ var
-      _ -> searchForConstraints eo typ var
-  | otherwise = searchForConstraints eo typ var
+searchForConstraints eo@(ExtraTypeScriptOptions {..}) typ var
+  | (ConT name, VarT keyVar : trailing) <- splitAppsT typ
+  , keyVar == var
+  , name `L.elem` typeFamiliesToMapToTypeScript = lift (reify name) >>= \case
+      FamilyI (ClosedTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _) _) _ ->
+        tell [GenericInfo var (TypeFamilyKey (typeFamilyTableName typeFamilyName trailing))]
+      FamilyI (OpenTypeFamilyD (TypeFamilyHead typeFamilyName _ _ _)) _ ->
+        tell [GenericInfo var (TypeFamilyKey (typeFamilyTableName typeFamilyName trailing))]
+      _ -> return ()
 searchForConstraints eo (AppT typ1 typ2) var = searchForConstraints eo typ1 var >> searchForConstraints eo typ2 var
 searchForConstraints eo (SigT typ _) var = searchForConstraints eo typ var
 searchForConstraints eo (InfixT typ1 _ typ2) var = searchForConstraints eo typ1 var >> searchForConstraints eo typ2 var
@@ -56,6 +56,6 @@
 hasFreeTypeVariable _ = False
 
 unifyGenericVariable :: [GenericInfo] -> String
-unifyGenericVariable genericInfos = case [nameBase name | GenericInfo _ (TypeFamilyKey name) <- genericInfos] of
+unifyGenericVariable genericInfos = case L.nub [name | GenericInfo _ (TypeFamilyKey name) <- genericInfos] of
   [] -> ""
-  names -> " extends keyof " <> (L.intercalate " & " names)
+  names -> " extends " <> (L.intercalate " & " ["keyof " <> n | n <- names])

src/Data/Aeson/TypeScript/Types.hs

@@ -250,5 +250,5 @@ data GenericInfo = GenericInfo Name GenericInfoExtra
   deriving Show
 
 data GenericInfoExtra = NormalStar
-                      | TypeFamilyKey Name
+                      | TypeFamilyKey String
   deriving Show

src/Data/Aeson/TypeScript/Util.hs

@@ -1,8 +1,8 @@
 {-# LANGUAGE CPP #-}
-{-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
 
 module Data.Aeson.TypeScript.Util where
 
@@ -169,6 +169,31 @@
     [x] -> mapType genericVariables x
     xs -> applyToArgsT (ConT $ tupleTypeName (L.length xs)) (fmap (mapType genericVariables) xs)
 
+-- | Split a type application into its head and arguments.
+splitAppsT :: Type -> (Type, [Type])
+splitAppsT = go []
+  where
+    go acc (AppT f x) = go (x : acc) f
+    go acc (SigT t _) = go acc t
+    go acc (ParensT t) = go acc t
+    go acc t = (t, acc)
+
+-- | Name for the lookup interface a mapped type family produces
+typeFamilyTableName :: Name -> [Type] -> String
+typeFamilyTableName famName trailingArgs = nameBase famName <> concatMap typeNameFragment trailingArgs
+  where
+    typeNameFragment :: Type -> String
+    typeNameFragment (ConT n) = nameBase n
+    typeNameFragment (VarT n) = nameBase n
+    typeNameFragment (PromotedT n) = nameBase n
+    typeNameFragment (SigT t _) = typeNameFragment t
+    typeNameFragment (ParensT t) = typeNameFragment t
+    typeNameFragment (AppT f x) = typeNameFragment f <> typeNameFragment x
+    typeNameFragment ListT = "List"
+    typeNameFragment (TupleT 0) = "Unit"
+    typeNameFragment (TupleT n) = "Tuple" <> show n
+    typeNameFragment _ = ""
+
 mapType :: [(Name, (Suffix, Var))] -> Type -> Type
 mapType g x@(VarT name) = tryPromote x g name
 mapType g x@(ConT name) = tryPromote x g name
@@ -224,7 +249,7 @@
 nothingOnFail action = recover (return Nothing) (Just <$> action)
 
 tryGetDoc :: (String -> String) -> Name -> Q Exp
 tryGetDoc haddockModifier n = do
 #if MIN_VERSION_template_haskell(2,18,0)
   maybeDoc <- nothingOnFail (getDoc (DeclDoc n)) >>= \case
     Just (Just doc) -> return $ Just $ Just $ haddockModifier doc

test/Spec.hs

@@ -8,7 +8,7 @@
 import qualified ClosedTypeFamilies
 import qualified Formatting
 import qualified Generic
 import qualified GetDoc
 import qualified HigherKind
 import qualified MaybeTuples
 
@@ -24,6 +24,7 @@
 import qualified UntaggedNoTagSingleConstructors
 import qualified UntaggedTagSingleConstructors
 import qualified UnwrapUnaryRecords
+import qualified UseSiteFamilyImage
 
 
 main :: IO ()
@@ -50,3 +51,4 @@
   UntaggedNoTagSingleConstructors.tests
   UntaggedTagSingleConstructors.tests
   UnwrapUnaryRecords.allTests
+  UseSiteFamilyImage.tests