{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}

module Constrained.Env where

import Constrained.Core (
  Var (..),
 )
import Constrained.GenT (
  MonadGenError,
  genError,
 )
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Typeable
import Prettyprinter

-- | Typed environments for mapping `Var a` to `a`
newtype Env = Env {Env -> Map EnvKey EnvValue
unEnv :: Map EnvKey EnvValue}
  deriving newtype (NonEmpty Env -> Env
Env -> Env -> Env
(Env -> Env -> Env)
-> (NonEmpty Env -> Env)
-> (forall b. Integral b => b -> Env -> Env)
-> Semigroup Env
forall b. Integral b => b -> Env -> Env
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
$c<> :: Env -> Env -> Env
<> :: Env -> Env -> Env
$csconcat :: NonEmpty Env -> Env
sconcat :: NonEmpty Env -> Env
$cstimes :: forall b. Integral b => b -> Env -> Env
stimes :: forall b. Integral b => b -> Env -> Env
Semigroup, Semigroup Env
Env
Semigroup Env =>
Env -> (Env -> Env -> Env) -> ([Env] -> Env) -> Monoid Env
[Env] -> Env
Env -> Env -> Env
forall a.
Semigroup a =>
a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
$cmempty :: Env
mempty :: Env
$cmappend :: Env -> Env -> Env
mappend :: Env -> Env -> Env
$cmconcat :: [Env] -> Env
mconcat :: [Env] -> Env
Monoid)
  deriving stock (Int -> Env -> ShowS
[Env] -> ShowS
Env -> String
(Int -> Env -> ShowS)
-> (Env -> String) -> ([Env] -> ShowS) -> Show Env
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Env -> ShowS
showsPrec :: Int -> Env -> ShowS
$cshow :: Env -> String
show :: Env -> String
$cshowList :: [Env] -> ShowS
showList :: [Env] -> ShowS
Show)

data EnvValue where
  EnvValue :: (Typeable a, Show a) => !a -> EnvValue

deriving instance Show EnvValue

data EnvKey where
  EnvKey :: !(Var a) -> EnvKey

instance Eq EnvKey where
  EnvKey Var a
v == :: EnvKey -> EnvKey -> Bool
== EnvKey Var a
v' = Var a -> Int
forall a. Var a -> Int
nameOf Var a
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Var a -> Int
forall a. Var a -> Int
nameOf Var a
v'

instance Ord EnvKey where
  compare :: EnvKey -> EnvKey -> Ordering
compare (EnvKey Var a
v) (EnvKey Var a
v') = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Var a -> Int
forall a. Var a -> Int
nameOf Var a
v) (Var a -> Int
forall a. Var a -> Int
nameOf Var a
v')

instance Show EnvKey where
  show :: EnvKey -> String
show (EnvKey Var a
var) = Var a -> String
forall a. Show a => a -> String
show Var a
var

extendEnv :: (Typeable a, Show a) => Var a -> a -> Env -> Env
extendEnv :: forall a. (Typeable a, Show a) => Var a -> a -> Env -> Env
extendEnv Var a
v a
a (Env Map EnvKey EnvValue
m) = Map EnvKey EnvValue -> Env
Env (Map EnvKey EnvValue -> Env) -> Map EnvKey EnvValue -> Env
forall a b. (a -> b) -> a -> b
$ EnvKey -> EnvValue -> Map EnvKey EnvValue -> Map EnvKey EnvValue
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Var a -> EnvKey
forall a. Var a -> EnvKey
EnvKey Var a
v) (a -> EnvValue
forall a. (Typeable a, Show a) => a -> EnvValue
EnvValue a
a) Map EnvKey EnvValue
m

removeVar :: Var a -> Env -> Env
removeVar :: forall a. Var a -> Env -> Env
removeVar Var a
v (Env Map EnvKey EnvValue
m) = Map EnvKey EnvValue -> Env
Env (Map EnvKey EnvValue -> Env) -> Map EnvKey EnvValue -> Env
forall a b. (a -> b) -> a -> b
$ EnvKey -> Map EnvKey EnvValue -> Map EnvKey EnvValue
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Var a -> EnvKey
forall a. Var a -> EnvKey
EnvKey Var a
v) Map EnvKey EnvValue
m

singletonEnv :: (Typeable a, Show a) => Var a -> a -> Env
singletonEnv :: forall a. (Typeable a, Show a) => Var a -> a -> Env
singletonEnv Var a
v a
a = Map EnvKey EnvValue -> Env
Env (Map EnvKey EnvValue -> Env) -> Map EnvKey EnvValue -> Env
forall a b. (a -> b) -> a -> b
$ EnvKey -> EnvValue -> Map EnvKey EnvValue
forall k a. k -> a -> Map k a
Map.singleton (Var a -> EnvKey
forall a. Var a -> EnvKey
EnvKey Var a
v) (a -> EnvValue
forall a. (Typeable a, Show a) => a -> EnvValue
EnvValue a
a)

lookupEnv :: Typeable a => Env -> Var a -> Maybe a
lookupEnv :: forall a. Typeable a => Env -> Var a -> Maybe a
lookupEnv (Env Map EnvKey EnvValue
m) Var a
v = do
  EnvValue a
val <- EnvKey -> Map EnvKey EnvValue -> Maybe EnvValue
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Var a -> EnvKey
forall a. Var a -> EnvKey
EnvKey Var a
v) Map EnvKey EnvValue
m
  a -> Maybe a
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a
val

findEnv :: (Typeable a, MonadGenError m) => Env -> Var a -> m a
findEnv :: forall a (m :: * -> *).
(Typeable a, MonadGenError m) =>
Env -> Var a -> m a
findEnv Env
env Var a
var = do
  case Env -> Var a -> Maybe a
forall a. Typeable a => Env -> Var a -> Maybe a
lookupEnv Env
env Var a
var of
    Just a
a -> a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
    Maybe a
Nothing -> String -> m a
forall (m :: * -> *) a. MonadGenError m => String -> m a
genError (String
"Couldn't find " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Var a -> String
forall a. Show a => a -> String
show Var a
var String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" in " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Env -> String
forall a. Show a => a -> String
show Env
env)

instance Pretty EnvValue where
  pretty :: forall ann. EnvValue -> Doc ann
pretty (EnvValue a
x) = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (String -> Doc ann) -> String -> Doc ann
forall a b. (a -> b) -> a -> b
$ Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
80 (a -> String
forall a. Show a => a -> String
show a
x)

instance Pretty EnvKey where
  pretty :: forall ann. EnvKey -> Doc ann
pretty (EnvKey Var a
x) = Var a -> Doc ann
forall a ann. Show a => a -> Doc ann
viaShow Var a
x

instance Pretty Env where
  pretty :: forall ann. Env -> Doc ann
pretty (Env Map EnvKey EnvValue
m) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep (Doc ann
"Env" Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
: (((EnvKey, EnvValue) -> Doc ann)
-> [(EnvKey, EnvValue)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (EnvKey, EnvValue) -> Doc ann
forall {a} {a} {ann}. (Pretty a, Pretty a) => (a, a) -> Doc ann
f (Map EnvKey EnvValue -> [(EnvKey, EnvValue)]
forall k a. Map k a -> [(k, a)]
Map.toList Map EnvKey EnvValue
m)))
    where
      f :: (a, a) -> Doc ann
f (a
k, a
v) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep [a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
k, Doc ann
"->", a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v]