Add support for type variables

chiroptical · Feb 5, 2024 · 00bae1e · 00bae1e
1 parent 7b71439
commit 00bae1e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/src/DeriveHasField.hs b/src/DeriveHasField.hs
@@ -12,7 +12,6 @@ import Control.Monad
 import Data.Char (toLower)
 import Data.Foldable as Foldable
 import Data.Traversable (for)
-import Debug.Trace
 import GHC.Records
 import Language.Haskell.TH
 import Language.Haskell.TH.Datatype
@@ -24,7 +23,7 @@ makeDeriveHasField :: (String -> String) -> DatatypeInfo -> DecsQ
 makeDeriveHasField fieldModifier datatypeInfo = do
   -- We do not support sum of product types
   constructorInfo <- case datatypeInfo.datatypeCons of
-    [info] -> trace (show datatypeInfo) $ pure info
+    [info] -> pure info
     _ -> fail "deriveHasField: only supports product types with a single data constructor"
 
   -- We only support data and newtype declarations
@@ -50,12 +49,16 @@ makeDeriveHasField fieldModifier datatypeInfo = do
         wantedFieldName = lowerFirst $ fieldModifier currentFieldName
         litTCurrentField = litT $ strTyLit currentFieldName
         litTFieldWanted = litT $ strTyLit wantedFieldName
-        parentTypeConstructor = conT datatypeInfo.datatypeName
+        parentType =
+          foldl'
+            (\acc var -> appT acc (varT $ tyVarBndrToName var))
+            (conT datatypeInfo.datatypeName)
+            datatypeInfo.datatypeVars
      in if currentFieldName == wantedFieldName
           then fail "deriveHasField: after applying fieldModifier, field didn't change"
           else
             [d|
-              instance HasField $litTFieldWanted $parentTypeConstructor $(pure ty) where
+              instance HasField $litTFieldWanted $parentType $(pure ty) where
                 getField = $(appTypeE (varE $ mkName "getField") litTCurrentField)
               |]
   pure $ Foldable.concat decs
@@ -64,3 +67,8 @@ lowerFirst :: String -> String
 lowerFirst = \case
   [] -> []
   (x : xs) -> toLower x : xs
+
+tyVarBndrToName :: TyVarBndr flag -> Name
+tyVarBndrToName = \case
+  PlainTV name _ -> name
+  KindedTV name _ _ -> name
diff --git a/test/DeriveHasFieldSpec.hs b/test/DeriveHasFieldSpec.hs
@@ -26,18 +26,18 @@ someType =
     , someTypeSomeEitherField = Right 0
     }
 
-data OtherType a = OtherType
+data OtherType a b = OtherType
   { otherTypeField :: Maybe a
-  , otherTypeOtherField :: Maybe a
+  , otherTypeOtherField :: Either a b
   }
 
 deriveHasFieldWith (dropPrefix "otherType") ''OtherType
 
-otherType :: OtherType Int
+otherType :: OtherType Int String
 otherType =
   OtherType
     { otherTypeField = Just 0
-    , otherTypeOtherField = Nothing
+    , otherTypeOtherField = Right "hello"
     }
 
 spec :: Spec
@@ -50,4 +50,4 @@ spec = do
       someType.someEitherField `shouldBe` Right 0
     it "compiles and gets the right field" $ do
       otherType.field `shouldBe` Just 0
-      otherType.otherField `shouldBe` Nothing
+      otherType.otherField `shouldBe` Right "hello"