{-# LANGUAGE UndecidableInstances #-} -- Eq/Ord/Show deriving

module Momo.Env
  ( Env
  , fromSignature
  , addSignature

  , addValue
  , findValue
  , addType
  , findType
  , addModule
  , findModule

  , Binding(..)
  , addSpec

  , find
  , FindError(..)
  ) where

import Control.Monad.Catch (Exception, MonadThrow(throwM))
import Data.Map.Strict qualified as Map
import Data.Text (Text)

import Momo.CoreSyntax qualified as Core
import Momo.Ident (Ident)
import Momo.Ident qualified as Ident
import Momo.ModSyntax qualified as Mod
import Momo.Path (Path)
import Momo.Path qualified as Path
import Momo.Subst (Subst)

data Binding term
  = Value (Core.Val term)
  | Type (Mod.TypeDecl term)
  | Module (Mod.ModType term)

deriving instance Core.EqTerm term => Eq (Binding term)
deriving instance Core.OrdTerm term => Ord (Binding term)
deriving instance Core.ShowTerm term => Show (Binding term)

type Env term = Ident.Table (Binding term)

addSignature :: [Mod.Specification term] -> Env term -> Env term
addSignature :: forall {k} (term :: k).
[Specification term] -> Env term -> Env term
addSignature [Specification term]
specs Env term
env = (Specification term -> Env term -> Env term)
-> Env term -> [Specification term] -> Env term
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Specification term -> Env term -> Env term
forall {k} (term :: k). Specification term -> Env term -> Env term
addSpec Env term
env [Specification term]
specs

fromSignature :: [Mod.Specification term] -> Env term
fromSignature :: forall {k} (term :: k). [Specification term] -> Env term
fromSignature [Specification term]
specs = [Specification term] -> Env term -> Env term
forall {k} (term :: k).
[Specification term] -> Env term -> Env term
addSignature [Specification term]
specs Env term
forall a. Monoid a => a
mempty

addSpec
  :: Mod.Specification term
  -> Env term
  -> Env term
addSpec :: forall {k} (term :: k). Specification term -> Env term -> Env term
addSpec = \case
  Mod.ValueSig  Ident
ident Val term
vty  -> Ident -> Val term -> Env term -> Env term
forall {k} (term :: k). Ident -> Val term -> Env term -> Env term
addValue Ident
ident Val term
vty
  Mod.TypeSig   Ident
ident TypeDecl term
decl -> Ident -> TypeDecl term -> Env term -> Env term
forall {k} (term :: k).
Ident -> TypeDecl term -> Env term -> Env term
addType Ident
ident TypeDecl term
decl
  Mod.ModuleSig Ident
ident ModType term
mty  -> Ident -> ModType term -> Env term -> Env term
forall {k} (term :: k).
Ident -> ModType term -> Env term -> Env term
addModule Ident
ident ModType term
mty

addValue :: Ident -> Core.Val term -> Env term -> Env term
addValue :: forall {k} (term :: k). Ident -> Val term -> Env term -> Env term
addValue Ident
ident = Ident -> Binding term -> Env term -> Env term
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Ident
ident (Binding term -> Env term -> Env term)
-> (Val term -> Binding term) -> Val term -> Env term -> Env term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val term -> Binding term
forall {k} (term :: k). Val term -> Binding term
Value

addType :: Ident -> Mod.TypeDecl term -> Env term -> Env term
addType :: forall {k} (term :: k).
Ident -> TypeDecl term -> Env term -> Env term
addType Ident
ident = Ident
-> Binding term
-> Map Ident (Binding term)
-> Map Ident (Binding term)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Ident
ident (Binding term
 -> Map Ident (Binding term) -> Map Ident (Binding term))
-> (TypeDecl term -> Binding term)
-> TypeDecl term
-> Map Ident (Binding term)
-> Map Ident (Binding term)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeDecl term -> Binding term
forall {k} (term :: k). TypeDecl term -> Binding term
Type

addModule :: Ident -> Mod.ModType term -> Env term -> Env term
addModule :: forall {k} (term :: k).
Ident -> ModType term -> Env term -> Env term
addModule Ident
ident = Ident
-> Binding term
-> Map Ident (Binding term)
-> Map Ident (Binding term)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Ident
ident (Binding term
 -> Map Ident (Binding term) -> Map Ident (Binding term))
-> (ModType term -> Binding term)
-> ModType term
-> Map Ident (Binding term)
-> Map Ident (Binding term)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModType term -> Binding term
forall {k} (term :: k). ModType term -> Binding term
Module

data FindError
  = ValueFieldExpected Path
  | TypeFieldExpected Path
  | ModuleFieldExpected Path
  | BindingNotFound Ident
  | StructureExpected Path
  | StructureFieldNotFound Path Text
  deriving (FindError -> FindError -> Bool
(FindError -> FindError -> Bool)
-> (FindError -> FindError -> Bool) -> Eq FindError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FindError -> FindError -> Bool
$c/= :: FindError -> FindError -> Bool
== :: FindError -> FindError -> Bool
$c== :: FindError -> FindError -> Bool
Eq, Eq FindError
Eq FindError
-> (FindError -> FindError -> Ordering)
-> (FindError -> FindError -> Bool)
-> (FindError -> FindError -> Bool)
-> (FindError -> FindError -> Bool)
-> (FindError -> FindError -> Bool)
-> (FindError -> FindError -> FindError)
-> (FindError -> FindError -> FindError)
-> Ord FindError
FindError -> FindError -> Bool
FindError -> FindError -> Ordering
FindError -> FindError -> FindError
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: FindError -> FindError -> FindError
$cmin :: FindError -> FindError -> FindError
max :: FindError -> FindError -> FindError
$cmax :: FindError -> FindError -> FindError
>= :: FindError -> FindError -> Bool
$c>= :: FindError -> FindError -> Bool
> :: FindError -> FindError -> Bool
$c> :: FindError -> FindError -> Bool
<= :: FindError -> FindError -> Bool
$c<= :: FindError -> FindError -> Bool
< :: FindError -> FindError -> Bool
$c< :: FindError -> FindError -> Bool
compare :: FindError -> FindError -> Ordering
$ccompare :: FindError -> FindError -> Ordering
Ord, Int -> FindError -> ShowS
[FindError] -> ShowS
FindError -> String
(Int -> FindError -> ShowS)
-> (FindError -> String)
-> ([FindError] -> ShowS)
-> Show FindError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FindError] -> ShowS
$cshowList :: [FindError] -> ShowS
show :: FindError -> String
$cshow :: FindError -> String
showsPrec :: Int -> FindError -> ShowS
$cshowsPrec :: Int -> FindError -> ShowS
Show)

instance Exception FindError

findValue
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Core.Val term)
findValue :: forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (Val term)
findValue Path
path Env term
env =
  Path -> Env term -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (Binding term)
find Path
path Env term
env m (Binding term) -> (Binding term -> m (Val term)) -> m (Val term)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Value Val term
val ->
      Val term -> m (Val term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Val term
val
    Binding term
_ ->
      FindError -> m (Val term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (Val term)) -> FindError -> m (Val term)
forall a b. (a -> b) -> a -> b
$ Path -> FindError
ValueFieldExpected Path
path

findType
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Mod.TypeDecl term)
findType :: forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (TypeDecl term)
findType Path
path Env term
env =
  Path -> Env term -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (Binding term)
find Path
path Env term
env m (Binding term)
-> (Binding term -> m (TypeDecl term)) -> m (TypeDecl term)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Type TypeDecl term
decl ->
      TypeDecl term -> m (TypeDecl term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeDecl term
decl
    Binding term
_ ->
      FindError -> m (TypeDecl term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (TypeDecl term)) -> FindError -> m (TypeDecl term)
forall a b. (a -> b) -> a -> b
$ Path -> FindError
TypeFieldExpected Path
path

findModule
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Mod.ModType term)
findModule :: forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (ModType term)
findModule Path
path Env term
env =
  Path -> Env term -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (Binding term)
find Path
path Env term
env m (Binding term)
-> (Binding term -> m (ModType term)) -> m (ModType term)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Module ModType term
mty ->
      ModType term -> m (ModType term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ModType term
mty
    Binding term
_ ->
      FindError -> m (ModType term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (ModType term)) -> FindError -> m (ModType term)
forall a b. (a -> b) -> a -> b
$ Path -> FindError
ModuleFieldExpected Path
path

find
  :: forall term m
  . ( Core.CoreSyntax term
    , MonadThrow m
    )
  => Path
  -> Env term
  -> m (Binding term)
find :: forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (Binding term)
find Path
queryPath Env term
env =
  case Path
queryPath of
    Path.Ident{Ident
ident :: Path -> Ident
ident :: Ident
ident} ->
      case Ident -> Env term -> Maybe (Binding term)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Ident
ident Env term
env of
        Maybe (Binding term)
Nothing ->
          FindError -> m (Binding term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (Binding term)) -> FindError -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ Ident -> FindError
BindingNotFound Ident
ident
        Just Binding term
binding ->
          Binding term -> m (Binding term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Binding term
binding

    Path.Dot{Path
path :: Path -> Path
path :: Path
path, Text
field :: Path -> Text
field :: Text
field} ->
      forall (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (ModType term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Env term -> m (ModType term)
findModule @term Path
path Env term
env m (ModType term)
-> (ModType term -> m (Binding term)) -> m (Binding term)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Mod.Signature [Specification term]
sg ->
          Path -> Text -> Subst -> [Specification term] -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Text -> Subst -> [Specification term] -> m (Binding term)
findField Path
path Text
field Subst
forall a. Monoid a => a
mempty [Specification term]
sg
        Mod.FunctorType{} ->
          -- XXX: non-instantiated functor
          FindError -> m (Binding term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (Binding term)) -> FindError -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ Path -> FindError
StructureExpected Path
path

findField
  :: forall term m
  .  ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Text
  -> Subst
  -> [Mod.Specification term]
  -> m (Binding term)
findField :: forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Text -> Subst -> [Specification term] -> m (Binding term)
findField Path
path Text
field Subst
subst = \case
  [] ->
    FindError -> m (Binding term)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (FindError -> m (Binding term)) -> FindError -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ Path -> Text -> FindError
StructureFieldNotFound Path
path Text
field

  Mod.ValueSig Ident
ident Val term
vty : [Specification term]
remaining ->
    if Ident
ident.name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
field then
      Binding term -> m (Binding term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Binding term -> m (Binding term))
-> Binding term -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ Val term -> Binding term
forall {k} (term :: k). Val term -> Binding term
Value (forall (term :: k).
CoreSyntax term =>
Val term -> Subst -> Val term
forall {k} (term :: k).
CoreSyntax term =>
Val term -> Subst -> Val term
Core.substVal @term Val term
vty Subst
subst)
    else
      Path -> Text -> Subst -> [Specification term] -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Text -> Subst -> [Specification term] -> m (Binding term)
findField Path
path Text
field Subst
subst [Specification term]
remaining

  Mod.TypeSig Ident
ident TypeDecl term
decl : [Specification term]
remaining ->
    if Ident
ident.name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
field then
      Binding term -> m (Binding term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Binding term -> m (Binding term))
-> Binding term -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ TypeDecl term -> Binding term
forall {k} (term :: k). TypeDecl term -> Binding term
Type (TypeDecl term -> Subst -> TypeDecl term
forall {k} (term :: k).
CoreSyntax term =>
TypeDecl term -> Subst -> TypeDecl term
Mod.substTypeDecl TypeDecl term
decl Subst
subst)
    else
      let
        subst' :: Subst
subst' =
          Ident -> Path -> Subst -> Subst
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert
            Ident
ident
            (Path -> Text -> Path
Path.Dot Path
path Ident
ident.name)
            Subst
subst
      in
        Path -> Text -> Subst -> [Specification term] -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Text -> Subst -> [Specification term] -> m (Binding term)
findField Path
path Text
field Subst
subst' [Specification term]
remaining

  Mod.ModuleSig Ident
ident ModType term
mty : [Specification term]
remaining ->
    if Ident
ident.name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
field then
      Binding term -> m (Binding term)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Binding term -> m (Binding term))
-> Binding term -> m (Binding term)
forall a b. (a -> b) -> a -> b
$ ModType term -> Binding term
forall {k} (term :: k). ModType term -> Binding term
Module (ModType term -> Subst -> ModType term
forall {k} (term :: k).
CoreSyntax term =>
ModType term -> Subst -> ModType term
Mod.substModType ModType term
mty Subst
subst)
    else
      let
        subst' :: Subst
subst' =
          Ident -> Path -> Subst -> Subst
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert
            Ident
ident
            (Path -> Text -> Path
Path.Dot Path
path Ident
ident.name)
            Subst
subst
      in
        Path -> Text -> Subst -> [Specification term] -> m (Binding term)
forall {k} (term :: k) (m :: * -> *).
(CoreSyntax term, MonadThrow m) =>
Path -> Text -> Subst -> [Specification term] -> m (Binding term)
findField Path
path Text
field Subst
subst' [Specification term]
remaining