Add OverloadedLabels support for positional lenses

kcsongor · Apr 18, 2023 · f7ef335 · f7ef335
1 parent 4b34e49
commit f7ef335
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 6 deletions.
diff --git a/generic-lens/src/Data/Generics/Labels.hs b/generic-lens/src/Data/Generics/Labels.hs
@@ -1,10 +1,14 @@
 {-# LANGUAGE PackageImports #-}
 {-# LANGUAGE AllowAmbiguousTypes    #-}
 {-# LANGUAGE ConstraintKinds        #-}
+{-# LANGUAGE CPP                    #-}
 {-# LANGUAGE DataKinds              #-}
 {-# LANGUAGE FlexibleInstances      #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE MultiParamTypeClasses  #-}
+#if MIN_VERSION_base(4,12,0)
+{-# LANGUAGE NoStarIsType           #-}
+#endif
 {-# LANGUAGE PolyKinds              #-}
 {-# LANGUAGE ScopedTypeVariables    #-}
 {-# LANGUAGE TypeApplications       #-}
@@ -107,15 +111,18 @@ instance {-# INCOHERENT #-} AsConstructor name s t a b => Constructor name s t a
 instance {-# INCOHERENT #-} AsConstructor' name s a => Constructor name s s a a where
   constructorPrism = _Ctor' @name
 
-data LabelType = FieldType | LegacyConstrType | ConstrType
+data LabelType = FieldType | LegacyConstrType | ConstrType | PositionType
 
 type family ClassifyLabel (name :: Symbol) :: LabelType where
   ClassifyLabel name =
-    If (CmpSymbol "_@" name == 'LT && CmpSymbol "_[" name == 'GT)
-      'LegacyConstrType
-      ( If (CmpSymbol "@" name == 'LT && CmpSymbol "[" name == 'GT)
-          'ConstrType
-          'FieldType
+    If (StartsWithDigit name)
+      'PositionType
+      ( If (CmpSymbol "_@" name == 'LT && CmpSymbol "_[" name == 'GT)
+          'LegacyConstrType
+          ( If (CmpSymbol "@" name == 'LT && CmpSymbol "[" name == 'GT)
+              'ConstrType
+              'FieldType
+          )
       )
 
 instance ( labelType ~ ClassifyLabel name
@@ -144,3 +151,35 @@ instance ( Applicative f, Choice p, Constructor name s t a b
 instance ( Applicative f, Choice p, Constructor name s t a b
          ) => IsLabelHelper 'ConstrType name p f s t a b where
   labelOutput = constructorPrism @name
+
+#if MIN_VERSION_base(4,18,0)
+type StartsWithDigit name =
+  CmpSymbol "/" name == 'LT && CmpSymbol ":" name == 'GT
+
+class Position (i :: Nat) s t a b | s i -> a, t i -> b, s i b -> t, t i a -> s where
+  positionLens :: Lens s t a b
+
+instance {-# INCOHERENT #-} HasPosition i s t a b => Position i s t a b where
+  positionLens = position @i
+
+instance {-# INCOHERENT #-} HasPosition' i s a => Position i s s a a where
+  positionLens = position' @i
+
+instance ( Functor f, Position i s t a b, i ~ ParseNat name
+         ) => IsLabelHelper 'PositionType name (->) f s t a b where
+  labelOutput = positionLens @i
+
+type ParseNat name = ParseNat' 0 (UnconsSymbol name)
+
+type family ParseNat' acc m where
+  ParseNat' acc ('Just '(hd, tl)) =
+    ParseNat' (10 * acc + DigitToNat hd) (UnconsSymbol tl)
+  ParseNat' acc 'Nothing = acc
+
+type DigitToNat c =
+  If ('0' <=? c && c <=? '9')
+    (CharToNat c - CharToNat '0')
+    (TypeError ('Text "Invalid position number"))
+#else
+type StartsWithDigit name = 'False
+#endif
diff --git a/generic-lens/test/Spec.hs b/generic-lens/test/Spec.hs
@@ -271,6 +271,14 @@ tests = TestList $ map mkHUnitTest
   , (valLabel ^? #RecB . _1  ) ~=? Just 3
   , (valLabel ^? #RecB       ) ~=? Just (3, True)
   , (valLabel ^? #RecC       ) ~=? Nothing
+
+  , (valLabel ^. #1          ) ~=? 3
+  , let
+      i x = x :: Int
+      largeTuple  = (i 1, i 2, i 3, i 4, i 5, i 6, i 7, i 8, i 9, i 10, i 11, i 12, i 13, i 14, i 15)
+      largeTuple' = (i 1, i 2, i 3, i 4, i 5, i 6, i 7, i 8, i 9, i 10, i 11, i 13, i 13, i 14, i 15)
+    in
+      (largeTuple ^. #13, largeTuple & #12 +~ 1) ~=? (13, largeTuple')
 #endif
   , customTypesTest
   ]