Record of maps with compositional lookups and updates?

Some pseudocode:

data A = A
data B = B
data C = C
data D = D
data E = E
data F = F
data G = G

data A1 = A1 A B C
data A2 = A2 A
data A3 = A3 B C D
data A4 = A4 D E F
data A5 = A5 A1 A4 G

data Foo k = Foo
    {
        a1s :: Map.Map k A1,
        a2s :: Map.Map k A2,
        a3s :: Map.Map k A3,
        a4s :: Map.Map k A4,
        a5s :: Map.Map k A5,
--and my attempted solution would use
        -- e.g. [(A1, [(A, Unit), (B, Unit), (C, Unit)]), (A5, [(A1, Composite), (A4, Composite), (G, Unit) ]) ]
        componentMap :: Map.Map Type (Set Type),

        -- e.g. [(A, [A1, A2]), (A1, [A5, A1]) ]
        compositeMap :: Map.Map Type (Set Type)
    }

I'd like to construct some kind of data structure that looks like this. From here, I'd like to:

lookup :: Foo k -> k -> Either FailureReason v individual values; if we assume that we have populated maps, I'd like lookup foo a1 :: A1, but also transitive instances such as lookup foo a1 :: B or lookup foo a5 :: A1 (since this is shorthand for getA1fromA5 $ lookup foo a5) and lookup foo a5 :: B. I'm considering FailureReason = WrongType | NotPresent but this is probably excessive.
traversals over types such as an (indexed) traversal over (k, D) which should hit everything in A3, A4, A5

This could be implemented as a recursive search over componentMap and compositeMap..so long as they were populated by hand.

Since the above seem very much recursive, I feel this has a GHC.Generics solution. Possibly a lens/optics + generic-lens/generic-optics one?

Or is my solution one that doesn't need generics and its ilk, and is instead just writing some traversals and lenses to index into my structure?

The question then becomes: is this functionality already existing in some library? If not, is Generics the tool I'm looking for to implement it?

Solution

I'm assuming you don't actually want multiple maps here -- that is, a given key should correspond to exactly one value, not an A1 value from the a1s map and another A2 value from from the a2s map, etc.

Also, you haven't said what you want to do if there are multiple matches of a particular type within in a single value, for example if you have values of type:

data A6 = A6 A3 A4

and try to retrieve or traverse terms of type D. Below, I assume you want to retrieve and/or traverse only the "first" one encountered (e.g., the D in A3 only, ignoring the one in A4).

Anyway, you can do this with Data generics and some helpers from lens's Data.Data.Lens.

No special data type is needed. A plain Map is sufficient, with a sum type to represent the collection of values you want to store:

data Dat = D_A1 A1 | D_A2 A2 | D_A3 A3 | D_A4 A4 | D_A5 A5 deriving (Data)
type Foo k dat = Map k dat

To look up a (possibly deeply nested) value by key, we can use the biplate traversal from lens:

lookupFoo :: (Ord k, Typeable v, Data dat) => k -> Foo k dat -> Maybe v
lookupFoo k foo = do
  dat <- Map.lookup k foo
  firstOf biplate dat

Here, biplate recursively traverses all the subterms of type v in the term dat. The firstOf query returns the first matching term or Nothing if no terms are found. (The do block is running in the Maybe monad.)

To perform an indexed traversal, we can also use biplate, modified using taking 1 to traverse only the first match:

itraverseFoo :: (Applicative f, Typeable v, Data dat) => (k -> v -> f v) -> Foo k dat -> f (Foo k dat)
itraverseFoo f foo = Map.traverseWithKey f' foo
  where f' k dat = taking 1 biplate (f k) dat

The full code:

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Lens
import Control.Monad.Writer
import Data.Data
import Data.Data.Lens
import Data.Map (Map)
import qualified Data.Map as Map

data Dat = D_A1 A1 | D_A2 A2 | D_A3 A3 | D_A4 A4 | D_A5 A5 | D_A6 A6 deriving (Data)
type Foo k dat = Map k dat

lookupFoo :: (Ord k, Typeable v, Data dat) => k -> Foo k dat -> Maybe v
lookupFoo k foo = do
  dat <- Map.lookup k foo
  firstOf biplate dat

itraverseFoo :: (Applicative f, Typeable v, Data dat) => (k -> v -> f v) -> Foo k dat -> f (Foo k dat)
itraverseFoo f foo = Map.traverseWithKey f' foo
  where f' k dat = taking 1 biplate (f k) dat

data A = A deriving (Data, Show)
data B = B deriving (Data, Show)
data C = C deriving (Data, Show)
data D = D deriving (Data, Show)
data E = E deriving (Data, Show)
data F = F deriving (Data, Show)
data G = G deriving (Data, Show)

data A1 = A1 A B C deriving (Data, Show)
data A2 = A2 A deriving (Data, Show)
data A3 = A3 B C D deriving (Data, Show)
data A4 = A4 D E F deriving (Data, Show)
data A5 = A5 A1 A4 G deriving (Data, Show)
data A6 = A6 A3 A4 deriving (Data, Show)

foo :: Foo String Dat
foo = Map.fromList [ ("a1", D_A1 (A1 A B C))
                   , ("a3", D_A3 (A3 B C D))
                   , ("a4", D_A4 (A4 D E F))
                   , ("a5", D_A5 (A5 (A1 A B C) (A4 D E F) G))
                   , ("a6", D_A6 (A6 (A3 B C D) (A4 D E F)))
                   ]

find :: forall a k. k -> a -> Writer [k] a
find k a = tell [k] >> pure a

main = do
  print $ (lookupFoo "a1" foo :: Maybe A1)
  print $ (lookupFoo "a1" foo :: Maybe B)
  print $ (lookupFoo "a5" foo :: Maybe A1)
  print $ (lookupFoo "a5" foo :: Maybe B)
  print $ execWriter (itraverseFoo (find @D) foo)