content: {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE FunctionalDependencies #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# OPTIONS_GHC -Wno-orphans #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  module Spec.STLC where
  
  import Data.List
  import Data.String (IsString(..))
  
  import Control.Applicative
  import Control.Monad.Identity
  import Control.Monad.State
  
  import Refinery.ProofState
  import Refinery.Tactic
  
  import Test.Hspec
  
  -- Just a very simple version of Simply Typed Lambda Calculus,
  -- augmented with 'Hole' so that we can have
  -- incomplete extracts.
  data Term
    = Var String
    | Hole Int
    | Lam String Term
    | Pair Term Term
    deriving (Show, Eq)
  
  
  -- The type part of simply typed lambda calculus
  data Type
    = TVar String
    | Type :-> Type
    | TPair Type Type
    deriving (Show, Eq)
  
  data TacticState = TacticState { name :: Int, meta :: Int }
      deriving Show
  
  fresh :: MonadState TacticState m => m String
  fresh = do
      nm <- gets (show . name)
      modify (\s -> s { name = name s + 1 })
      pure nm
  
  infixr 4 :->
  
  instance IsString Type where
      fromString = TVar
  
  -- A judgement is just a context, along with a goal
  data Judgement = [(String, Type)] :- Type
    deriving (Show, Eq)
  
  instance MonadExtract Int Term String TacticState Identity where
    hole = do
      m <- gets meta
      modify $ \ts -> ts { meta = m + 1}
      pure (m, Hole m)
    unsolvableHole _ = do
      m <- gets meta
      modify $ \ts -> ts { meta = m + 1}
      pure (m, Hole m)
  
  instance MetaSubst Int Term where
      substMeta _ _ (Var s) = Var s
      substMeta i t1 (Hole i') = if i == i' then t1 else (Hole i')
      substMeta i t1 (Lam s body) = Lam s (substMeta i t1 body)
      substMeta i t1 (Pair l r) = Pair (substMeta i t1 l) (substMeta i t1 r)
  
  type T a = TacticT Judgement Term String TacticState Identity a
  
  pair :: T ()
  pair = rule $ \case
      (hys :- TPair a b) -> Pair <$> subgoal (hys :- a) <*> subgoal (hys :- b)
      _                  -> unsolvable "goal mismatch: Pair"
  
  lam :: T ()
  lam = rule $ \case
      (hys :- (a :-> b)) -> do
          nm <- fresh
          body <- subgoal $ ((nm, a) : hys) :- b
          pure $ Lam nm body
      _                  -> unsolvable "goal mismatch: Lam"
  
  assumption :: T ()
  assumption = rule $ \ (hys :- a) ->
    case find (\(_, ty) -> ty == a) hys of
      Just (x, _) -> pure $ Var x
      Nothing     -> unsolvable "goal mismatch: Assumption"
  
  auto :: T ()
  auto = do
      many_ lam
      choice [ pair >> auto
             , assumption
             ]
  
  refine :: T ()
  refine = do
      many_ lam
      try pair
  
  testHandlers :: T ()
  testHandlers = do
      handler (\err -> pure $ err ++ " Third")
      handler (\err -> pure $ err ++ " Second")
      failure "First"
  
  testHandlerAlt :: T ()
  testHandlerAlt = do
      handler (\err -> pure $ err ++ " Handled")
      (failure "Error1") <|> (failure "Error2")
  
  testReify :: T ()
  testReify = do
      lam
      lam
      reify pair $ \ (Proof _ _ goals) -> failure $ "Generated " <> show (length goals) <> " subgoals"
  
  testAttempt :: T ()
  testAttempt = do
      lam
      attempt lam (failure "Attempt Test Failed")
      failure $ "Attempt Test Succeeds"
  
  jdg :: Judgement
  jdg = ([] :- ("a" :-> "b" :-> (TPair "a" "b")))
  
  evalT :: T () -> Judgement -> Either [String] [Term]
  evalT t j = fmap (fmap pf_extract) $ runIdentity $ runTacticT t j (TacticState 0 0)
  
  stlcTests :: Spec
  stlcTests = do
      describe "Simply Typed Lambda Calculus" $ do
          it "auto synthesize a solution" $ (evalT auto jdg) `shouldBe` (Right [(Lam "0" $ Lam "1" $ Pair (Var "0") (Var "1"))])
          it "handler ordering is correct" $ (evalT testHandlers jdg) `shouldBe` (Left ["First Second Third"])
          it "handler works through alt" $ (evalT testHandlerAlt jdg) `shouldBe` (Left ["Error1 Handled","Error2 Handled"])
          it "reify gets the right subgoals" $ (evalT testReify jdg) `shouldBe` (Left ["Generated 2 subgoals"])
          it "attempt properly handles errors" $ (evalT testAttempt jdg) `shouldBe` (Left ["Attempt Test Succeeds"])
  