Getting fields from the "union" of two records. Inspired by section 3.1 or Hiromi Ishii's "Witness Me" functional pearl.

RecordUnionFields.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Getting fields from the "union" of two records.
--
-- Inspired by section 3.1 or Hiromi Ishii's "Witness Me — Constructive
-- Arguments Must Be Guided with Concrete Witness" functional pearl.
module Main where

import Data.Kind
import GHC.Generics (C1, D1, K1 (..), M1 (..), Rec0 (..), S1 (..))
import GHC.Generics qualified as G
import GHC.Records
import GHC.TypeLits

--
-- Structure of this file
--
-- - A type-level map with related type-level functions.
-- - Generics machinery for getting the type-level map of conventional records.
-- - The helper "record union" datatype and its HasField instance.
-- - An actual example.

--
-- A type-level binary search tree. The tree is unbalanced.
--
data Map symbol q
  = E
  | N (Map symbol q) symbol q (Map symbol q)
  deriving (Show, Eq)

type Insert :: Symbol -> q -> Map Symbol q -> Map Symbol q
type family Insert k v t where
  Insert k v E = N E k v E
  Insert k v (N left k' v' right) = Insert' (CmpSymbol k k') k v left k' v' right

type Insert' :: Ordering -> Symbol -> q -> Map Symbol q -> Symbol -> q -> Map Symbol q -> Map Symbol q
type family Insert' ordering k v left k' v' right where
  Insert' EQ k v left k' v' right = N left k v right -- overwrite
  Insert' LT k v left k' v' right = N (Insert k v left) k' v' right
  Insert' GT k v left k' v' right = N left k' v' (Insert k v right)

-- If a key exists in two maps, the value of the first map wins.
type Combine :: Map Symbol q -> Map Symbol q -> Map Symbol q
type family Combine twinner tloser where
  Combine E tloser = tloser
  Combine (N left k v right) tloser =
    Combine right (Insert k v (Combine left tloser))

type Find :: Symbol -> Map Symbol q -> q
type family Find k t where
  Find k E = TypeError (Text "The type " :<>: ShowType k :<>: Text " is not present in the record.")
  Find k (N left k' v' right) = Find' (CmpSymbol k k') k left v' right

type Find' :: Ordering -> Symbol -> Map Symbol q -> q -> Map Symbol q -> q
type family Find' ordering k left v right where
  Find' EQ _ _ v' _ = v'
  Find' LT k left _ _ = Find k left
  Find' GT k _ _ right = Find k right

-- Make all the "values" in the type level map equal to some argument "value".
type Overwrite :: q -> Map symbol z -> Map symbol q
type family Overwrite q t where
  Overwrite _ E = E
  Overwrite v (N left k' v' right) = N (Overwrite v left) k' v (Overwrite v right)

--
-- Helper machinery to get the type-level map of fields from a record, using generics.
--
type RecordCode :: Type -> Map Symbol Type
type family RecordCode r where
  RecordCode r = RecordCode' E (G.Rep r)

class ToRecord' (start :: Map Symbol Type) (g :: Type -> Type) where
  type RecordCode' start g :: Map Symbol Type

instance ToRecord' E fields => ToRecord' E (D1 meta (C1 metacons fields)) where
  type RecordCode' E (D1 meta (C1 metacons fields)) = RecordCode' E fields

instance
  ToRecord'
    start
    ( S1
        ( 'G.MetaSel
            ( 'Just k)
            unpackedness
            strictness
            laziness
        )
        (Rec0 v)
    )
  where
  type
    RecordCode'
      start
      ( S1
          ( 'G.MetaSel
              ( 'Just k)
              unpackedness
              strictness
              laziness
          )
          (Rec0 v)
      ) =
      Insert k v start

instance
  ( ToRecord' start t2,
    RecordCode' start t2 ~ middle,
    ToRecord' middle t1
  ) =>
  ToRecord' start (t1 G.:*: t2)
  where
  type RecordCode' start (t1 G.:*: t2) = RecordCode' (RecordCode' start t2) t1

--
-- The "record union" helper datatype and its HasField instance.
--
data RU r1 r2 = RU r1 r2 deriving (Show)

-- Should we search for a field in the first record, or in the second?
data Location
  = InFirst
  | InSecond

-- A version of HasField with an extra "where to search" parameter.
type UnionHasField :: Location -> Symbol -> Type -> Type -> Constraint
class UnionHasField location x r a | location x r -> a where
  getFieldU :: r -> a

instance HasField k r1 a => UnionHasField InFirst k (RU r1 r2) a where
  getFieldU (RU r1 _) = getField @k r1

instance HasField k r2 a => UnionHasField InSecond k (RU r1 r2) a where
  getFieldU (RU _ r2) = getField @k r2

-- The union of two records has a field when the type level map from field names from "where to search" has the field.
-- https://discourse.haskell.org/t/getting-fields-from-the-union-of-two-records/2296/2
instance
  (loc ~ Find k (Overwrite InFirst (RecordCode r1) `Combine` (Overwrite InSecond (RecordCode r2))),
  UnionHasField loc k (RU r1 r2) a)
  =>
  HasField k (RU r1 r2) a
  where
  getField r = getFieldU @loc @k r

--
-- A CONCRETE EXAMPLE
--
data Person = Person {name :: String, age :: Int} deriving (Show, G.Generic)

data Address = Address {street :: String, number :: Int} deriving (Show, G.Generic)

personAddress = RU (Person "John" 50) (Address "Rue del Percebe" 2)

main :: IO ()
main = do
  print $ getField @"name" personAddress
  print $ getField @"age" personAddress
  print $ getField @"street" personAddress
  print $ getField @"number" personAddress