{-# language GADTs, StandaloneDeriving #-}
{-# language LambdaCase #-}
{-# language MultiParamTypeClasses #-}
{-# language TemplateHaskell #-}
module Graph where

import Control.Applicative
import Data.Dependent.Map (DMap)
import Data.Dependent.Sum (ShowTag(..))
import Data.GADT.Compare.TH (deriveGEq, deriveGCompare)
import Data.GADT.Show.TH (deriveGShow)

import qualified Data.Dependent.Map as DMap

data ID a where
  ID_Expr :: Int -> ID Expr
  ID_Bound :: Int -> ID Bound
  ID_Binding :: Int -> ID Binding
deriving instance Eq (ID a)
deriving instance Ord (ID a)
deriving instance Show (ID a)

data Bound = Bound String
  deriving (Eq, Show)

data Binding = Binding String
  deriving (Eq, Show)

data Expr
  = Var Bound
  | App Expr Expr
  | Lam Binding Expr
  deriving (Eq, Show)

data NodeInfo a where
  BoundInfo ::
    -- value
    String ->
    NodeInfo Bound
  BindingInfo ::
    -- value
    String ->
    NodeInfo Binding
  VarInfo ::
    -- child
    ID Bound ->
    NodeInfo Expr
  AppInfo ::
    -- left child
    ID Expr ->
    -- right child
    ID Expr ->
    NodeInfo Expr
  LamInfo ::
    -- left child
    ID Binding ->
    -- right child
    ID Expr ->
    NodeInfo Expr
deriving instance Eq (NodeInfo a)
deriving instance Ord (NodeInfo a)
deriving instance Show (NodeInfo a)

deriveGEq ''ID
deriveGCompare ''ID
deriveGShow ''ID

instance ShowTag ID NodeInfo where
  showTaggedPrec ID_Expr{} = showsPrec
  showTaggedPrec ID_Bound{} = showsPrec
  showTaggedPrec ID_Binding{} = showsPrec

class Unbuild a where
  unbuild :: DMap ID NodeInfo -> a -> (ID a, DMap ID NodeInfo)

instance Unbuild Bound where
  unbuild m (Bound a) =
    let
      i' = (ID_Bound $ DMap.size m)
    in
      (i', DMap.insert i' (BoundInfo a) m)

instance Unbuild Binding where
  unbuild m (Binding a) =
    let
      i' = (ID_Binding $ DMap.size m)
    in
      (i', DMap.insert i' (BindingInfo a) m)

instance Unbuild Expr where
  unbuild m (Var a) =
    let
      (a', m') = unbuild m a
      i' = (ID_Expr $ DMap.size m')
    in
      (i', DMap.insert i' (VarInfo a') m')
  unbuild m (App a b) =
    let
      (a', m') = unbuild m a
      (b', m'') = unbuild m' b
      i' = (ID_Expr $ DMap.size m'')
    in
      (i', DMap.insert i' (AppInfo a' b') m'')
  unbuild m (Lam a b) =
    let
      (a', m') = unbuild m a
      (b', m'') = unbuild m' b
      i' = (ID_Expr $ DMap.size m'')
    in
      (i', DMap.insert i' (LamInfo a' b') m'')

rebuild :: DMap ID NodeInfo -> ID a -> Maybe a
rebuild m i =
  DMap.lookup i m >>=
  \case
    BoundInfo val -> pure $ Bound val
    BindingInfo val -> pure $ Binding val
    VarInfo val -> Var <$> rebuild m val
    AppInfo a b -> App <$> rebuild m a <*> rebuild m b
    LamInfo a b -> Lam <$> rebuild m a <*> rebuild m b