records-edsl-deriving-optics-0.1.0: Records/EDSL/Deriving/Optics.hs

module Records.EDSL.Deriving.Optics (deriveLabelOptics) where

import Data.List ((!!))
import Records.EDSL.Deriving.Type
import Records.EDSL.Description
import Language.Haskell.TH qualified as TH
import Optics
import Relude hiding (lines)

deriveLabelOptics :: Deriver
deriveLabelOptics = deriver #optics_LabelOptic \RecordDesc{constrName, fields, typeName} ->
  concat <$> ifor fields \i FieldDesc {nameText, type_} ->
    let qs = TH.conT typeName
        qtype = pure type_.hask.type_
        fieldCount = length fields
     in [d|
          instance (a ~ $qtype, b ~ $qtype) => LabelOptic $(TH.litT (TH.strTyLit (toString nameText))) A_Lens $qs $qs a b where
            {-# INLINE labelOptic #-}
            labelOptic = $(makeLensExp [(constrName, fieldCount, [i])])
          |]
  where
    -- Lifted from optics Optics.TH.Internal.Product
   
    -- | Build a lens clause that updates the field at the given index
    makeLensExp :: [(TH.Name, Int, [Int])] -> TH.ExpQ
    makeLensExp constrs = do
      f <- TH.newName "f"
      s <- TH.newName "s"
      TH.appsE
        [ TH.varE 'lensVL
        , TH.lamE [TH.varP f, TH.varP s] $
            TH.caseE
              (TH.varE s)
              [ makeLensMatch f conName fieldCount fields
              | (conName, fieldCount, fields) <- constrs
              ]
        ]

     -- | Make a lens match. Used for both lens and affine traversal generation.
    makeLensMatch :: TH.Name -> TH.Name -> Int -> [Int] -> TH.Q TH.Match
    makeLensMatch f conName fieldCount = \case
      [field] -> do
        xs <- newNames "x" fieldCount
        y <- TH.newName "y"
        let body =
              TH.appsE
                [ TH.varE 'fmap
                , TH.lamE [TH.varP y] . TH.appsE $
                    TH.conE conName : map TH.varE (set (ix field) y xs)
                , TH.appE (TH.varE f) . TH.varE $ xs !! field
                ]
        -- Con x_1 .. x_n -> fmap (\y_i -> Con x_1 .. y_i .. x_n) (f x_i)
        TH.match
          (TH.conP conName $ map TH.varP xs)
          (TH.normalB body)
          []
      _ -> error "Lens focuses on exactly one field"

    -- | Generate many new names from a given base name.
    newNames :: String {- ^ base name -} -> Int {- ^ count -} -> TH.Q [TH.Name]
    newNames base n = sequence [ TH.newName (base ++ show i) | i <- [1..n] ]