content: -- |
  -- Module      :  Cryptol.Utils.RecordMap
  -- Copyright   :  (c) 2020 Galois, Inc.
  -- License     :  BSD3
  -- Maintainer  :  cryptol@galois.com
  -- Stability   :  provisional
  -- Portability :  portable
  --
  -- This module implements an "order insensitive" datastructure for
  -- record types and values.  For most purposes, we want to deal with
  -- record fields in a canonical order; but for user interaction
  -- purposes, we generally want to display the fields in the order they
  -- were specified by the user (in source files, at the REPL, etc.).
  
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE Safe #-}
  
  module Cryptol.Utils.RecordMap
    ( RecordMap
    , displayOrder
    , canonicalFields
    , displayFields
    , recordElements
    , displayElements
    , fieldSet
    , recordFromFields
    , recordFromFieldsErr
    , recordFromFieldsWithDisplay
    , lookupField
    , adjustField
    , traverseRecordMap
    , mapWithFieldName
    , zipRecordsM
    , zipRecords
    , recordMapAccum
    ) where
  
  import           Control.DeepSeq
  import           Control.Monad.Except (ExceptT, MonadError(..), runExceptT)
  import           Control.Monad.Trans (MonadTrans(..))
  import           Data.Functor.Identity
  import           Data.Set (Set)
  import           Data.Map (Map)
  import qualified Data.Map.Strict as Map
  import qualified Data.Map.Merge.Strict as Map
  
  import Cryptol.Utils.Panic
  
  -- | An "order insensitive" datastructure.
  --   The fields can be accessed either according
  --   to a "canonical" order, or based on a
  --   "display" order, which matches the order
  --   in which the fields were originally specified.
  data RecordMap a b =
    RecordMap
    { recordMap :: !(Map a b)
    , _displayOrder :: [a]
    }
  -- RecordMap Invariant:
  --   The `displayOrder` field should contain exactly the
  --   same set of field names as the keys of `recordMap`.
  --   Moreover, each field name should occur at most once.
  
  instance (Ord a, Eq b) => Eq (RecordMap a b) where
    a == b = recordMap a == recordMap b
  
  instance (Ord a, Ord b) => Ord (RecordMap a b) where
    compare a b = compare (recordMap a) (recordMap b)
  
  instance (Show a, Ord a, Show b) => Show (RecordMap a b) where
    show = show . displayFields
  
  instance (NFData a, NFData b) => NFData (RecordMap a b) where
    rnf = rnf . canonicalFields
  
  
  -- | Return the fields in this record as a set.
  fieldSet :: Ord a => RecordMap a b -> Set a
  fieldSet r = Map.keysSet (recordMap r)
  
  -- | The order in which the fields originally appeared.
  displayOrder :: RecordMap a b -> [a]
  displayOrder r = _displayOrder r
  
  -- | Retrieve the elements of the record in canonical order
  --   of the field names
  recordElements :: RecordMap a b -> [b]
  recordElements = map snd . canonicalFields
  
  -- | Return a list of field/value pairs in canonical order.
  canonicalFields :: RecordMap a b -> [(a,b)]
  canonicalFields = Map.toList . recordMap
  
  -- | Retrieve the elements of the record in display order
  --   of the field names.
  displayElements :: (Show a, Ord a) => RecordMap a b -> [b]
  displayElements = map snd . displayFields
  
  -- | Return a list of field/value pairs in display order.
  displayFields :: (Show a, Ord a) => RecordMap a b -> [(a,b)]
  displayFields r = map find (displayOrder r)
    where
    find x =
      case Map.lookup x (recordMap r) of
        Just v -> (x, v)
        Nothing ->
           panic "displayFields"
                 ["Could not find field in recordMap " ++ show x
                 , "Display order: " ++ show (displayOrder r)
                 , "Canonical order:" ++ show (map fst (canonicalFields r))
                 ]
  
  -- | Generate a record from a list of field/value pairs.
  --   Precondition: each field identifier appears at most
  --   once in the given list.
  recordFromFields :: (Show a, Ord a) => [(a,b)] -> RecordMap a b
  recordFromFields xs =
    case recordFromFieldsErr xs of
      Left (x,_) ->
            panic "recordFromFields"
                  ["Repeated field value: " ++ show x]
      Right r -> r
  
  -- | Generate a record from a list of field/value pairs.
  --   If any field name is repeated, the first repeated name/value
  --   pair is returned.  Otherwise the constructed record is returned.
  recordFromFieldsErr :: (Show a, Ord a) => [(a,b)] -> Either (a,b) (RecordMap a b)
  recordFromFieldsErr xs0 = loop mempty xs0
    where
    loop m [] = Right (RecordMap m (map fst xs0))
    loop m ((x,v):xs)
      | Just _ <- Map.lookup x m = Left (x,v)
      | otherwise = loop (Map.insert x v m) xs
  
  -- | Generate a record from a list of field/value pairs,
  --   and also provide the "display" order for the fields directly.
  --   Precondition: each field identifier appears at most once in each
  --   list, and a field name appears in the display order iff it appears
  --   in the field list.
  recordFromFieldsWithDisplay :: (Show a, Ord a) => [a] -> [(a,b)] -> RecordMap a b
  recordFromFieldsWithDisplay dOrd fs = r { _displayOrder = dOrd }
    where
    r = recordFromFields fs
  
  -- | Lookup the value of a field
  lookupField :: Ord a => a -> RecordMap a b -> Maybe b
  lookupField x m = Map.lookup x (recordMap m)
  
  -- | Update the value of a field by applying the given function.
  --   If the field is not present in the record, return Nothing.
  adjustField :: forall a b. Ord a => a -> (b -> b) -> RecordMap a b -> Maybe (RecordMap a b)
  adjustField x f r = mkRec <$> Map.alterF f' x (recordMap r)
    where
    mkRec m = r{ recordMap = m }
  
    f' :: Maybe b -> Maybe (Maybe b)
    f' Nothing = Nothing
    f' (Just v) = Just (Just (f v))
  
  -- | Traverse the elements of the given record map in canonical
  --   order, applying the given action.
  traverseRecordMap :: Applicative t =>
    (a -> b -> t c) -> RecordMap a b -> t (RecordMap a c)
  traverseRecordMap f r =
    (\m -> RecordMap m (displayOrder r)) <$> Map.traverseWithKey f (recordMap r)
  
  -- | Apply the given function to each element of a record.
  mapWithFieldName :: (a -> b -> c) -> RecordMap a b -> RecordMap a c
  mapWithFieldName f = runIdentity . traverseRecordMap (\a b -> Identity (f a b))
  
  instance Functor (RecordMap a) where
    fmap f = mapWithFieldName (\_ -> f)
  
  instance Foldable (RecordMap a) where
    foldMap f = foldMap (f . snd) . canonicalFields
  
  instance Traversable (RecordMap a) where
    traverse f = traverseRecordMap (\_ -> f)
  
  -- | The function recordMapAccum threads an accumulating argument through
  --   the map in canonical order of fields.
  recordMapAccum :: (a -> b -> (a,c)) -> a -> RecordMap k b -> (a, RecordMap k c)
  recordMapAccum f z r = (a, r{ recordMap = m' })
    where
    (a, m') = Map.mapAccum f z (recordMap r)
  
  -- | Zip together the fields of two records using the provided action.
  --   If some field is present in one record, but not the other,
  --   an @Either a a@ will be returned, indicating which record is missing
  --   the field, and returning the name of the missing field.
  --
  --   The @displayOrder@ of the resulting record will be taken from the first
  --   argument (rather arbitrarily).
  zipRecordsM :: forall t a b c d. (Ord a, Monad t) =>
    (a -> b -> c -> t d) -> RecordMap a b -> RecordMap a c -> t (Either (Either a a) (RecordMap a d))
  zipRecordsM f r1 r2 = runExceptT (mkRec <$> zipMap (recordMap r1) (recordMap r2))
    where
    mkRec m = RecordMap m (displayOrder r1)
  
    zipMap :: Map a b -> Map a c -> ExceptT (Either a a) t (Map a d)
    zipMap = Map.mergeA missingLeft missingRight matched
    missingLeft  = Map.traverseMissing (\a _b -> throwError (Left a))
    missingRight = Map.traverseMissing (\a _c -> throwError (Right a))
    matched = Map.zipWithAMatched (\a b c -> lift (f a b c))
  
  -- | Pure version of `zipRecordsM`
  zipRecords :: forall a b c d. Ord a =>
    (a -> b -> c -> d) -> RecordMap a b -> RecordMap a c -> Either (Either a a) (RecordMap a d)
  zipRecords f r1 r2 = runIdentity (zipRecordsM (\a b c -> Identity (f a b c)) r1 r2)
  