packages feed

disco-0.2: src/Disco/Interpret/CESK.hs

{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fmax-pmcheck-models=200 #-}

-- |
-- Module      :  Disco.Interpret.CESK
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- CESK machine interpreter for Disco.
module Disco.Interpret.CESK (
  CESK,
  runCESK,
  step,
  eval,
  evalApp,
  runTest,
)
where

import Unbound.Generics.LocallyNameless (Bind, Name)

import Algebra.Graph
import qualified Algebra.Graph.AdjacencyMap as AdjMap
import Control.Arrow ((***), (>>>))
import Control.Monad ((>=>))
import Data.Bifunctor (first, second)
import Data.Functor (($>))
import Data.List (find)
import qualified Data.List.Infinite as InfList
import qualified Data.Map as M
import Data.Maybe (isJust)
import Data.Ratio
import qualified Data.Text as T
import Disco.AST.Core
import Disco.AST.Generic (
  Ellipsis (..),
  Side (..),
  selectSide,
 )
import Disco.AST.Typed (AProperty)
import Disco.Compile
import Disco.Context as Ctx
import Disco.Effects.Fresh
import Disco.Effects.Input
import Disco.Enumerate
import Disco.Error
import Disco.Names
import Disco.Property
import Disco.Syntax.Operators (BOp (..))
import Disco.Types hiding (V)
import Disco.Value
import Math.Combinatorics.Exact.Binomial (choose)
import Math.Combinatorics.Exact.Factorial (factorial)
import Math.NumberTheory.Primes (factorise, unPrime)
import Math.NumberTheory.Primes.Testing (isPrime)
import Math.OEIS (
  SearchStatus (SubSeq),
  extendSeq,
  lookupSeq,
  number,
 )
import Polysemy
import Polysemy.Error
import Polysemy.Random
import Polysemy.State
import qualified System.Random as R

------------------------------------------------------------
-- Utilities
------------------------------------------------------------

------------------------------------------------------------
-- Frames and continuations
------------------------------------------------------------

-- The CESK machine carries a current continuation explaining what to
-- do with the value of the currently focused expression, once it has
-- been fully evaluated.

-- | A continuation is just a stack of frames.
type Cont = [Frame]

-- | A frame represents a single step of the context, explaining what
--   to do with a value in that context (ultimately transforming it
--   into another value, which may in turn be handed to the next frame
--   in the continuation stack, and so on).
--
--   As an invariant, any 'Frame' containing unevaluated 'Core'
--   expressions must also carry an 'Env' in which to evaluate them.
data Frame
  = -- | Inject the value into a sum type.
    FInj Side
  | -- | Do a case analysis on the value.
    FCase Env (Bind (Name Core) Core) (Bind (Name Core) Core)
  | -- | Evaluate the right-hand value of a pair once we have finished
    --   evaluating the left-hand side.
    FPairR Env Core
  | -- | Put the value into the right-hand side of a pair together with
    --   this previously evaluated left-hand side.
    FPairL Value
  | -- | Project one or the other side of a pair.
    FProj Side
  | -- | Evaluate the argument of an application once we have finished
    --   evaluating the function.
    FArg Env Core
  | -- | Apply an evaluated function to this already-evaluated argument.
    FArgV Value
  | -- | Apply a previously evaluated function to the value.
    FApp Value
  | -- | Force evaluation of the contents of a memory cell.
    FForce
  | -- | Update the contents of a memory cell with its evaluation.
    FUpdate Int
  | -- | Record the results of a test.
    FTest TestVars Env
  | -- | Given the index of a memory cell and a function's arguments,
    --   memoize the results of a function.
    FMemo Int SimpleValue
  deriving (Show)

------------------------------------------------------------
-- The CESK machine
------------------------------------------------------------

-- | The CESK machine has two basic kinds of states.
data CESK
  = -- | The 'In' constructor represents the state when we are recursing
    --   "into" a term.  There is a currently focused expression which
    --   is to be evaluated in the given context.  Generally, evaluation
    --   proceeds by pattern-matching on the focused expression and
    --   either immediately turning it into a value (if it is simple),
    --   or focusing on a subexpression and pushing a new frame on the
    --   continuation stack indicating how to continue evaluating the
    --   whole expression once finished with the subexpression.
    In Core Env Cont
  | -- | The 'Out' constructor represents the state when we have
    --   completed evaluating an expression and are now on our way back
    --   "out" of the recursion.  Generally, evaluation proceeds by
    --   pattern-matching on the top frame of the continuation stack
    --   (and sometimes on the value as well), to see what is to be done
    --   with the value.
    Out Value Cont
  | -- | There is also an 'Up' constructor representing an exception
    --   that is propagating up the continuation stack.  Disco does
    --   not have user-level exceptions or try/catch blocks etc., but
    --   exceptions may be caught by test frames and turned into a
    --   test result rather than crashing the entire computation.
    Up EvalError Cont
  deriving (Show)

-- | Is the CESK machine in a final state?
isFinal :: CESK -> Maybe (Either EvalError Value)
isFinal (Up e []) = Just (Left e)
isFinal (Out v []) = Just (Right v)
isFinal _ = Nothing

-- | Run a CESK machine to completion.
runCESK :: Members '[Fresh, Random, State Mem] r => CESK -> Sem r (Either EvalError Value)
runCESK cesk = case isFinal cesk of
  Just res -> return res
  Nothing -> step cesk >>= runCESK

-- | Advance the CESK machine by one step.
step :: Members '[Fresh, Random, State Mem] r => CESK -> Sem r CESK
step cesk = case cesk of
  (In (CVar x) e k) -> case Ctx.lookup' x e of
    Nothing -> return $ Up (UnboundError x) k
    Just v -> return $ Out v k
  (In (CNum r) _ k) -> return $ Out (VNum r) k
  (In (CConst OMatchErr) _ k) -> return $ Up NonExhaustive k
  (In (CConst OEmptyGraph) _ k) -> return $ Out (VGraph empty) k
  (In (CConst op) _ k) -> return $ Out (VConst op) k
  (In (CInj s c) e k) -> return $ In c e (FInj s : k)
  (In (CCase c b1 b2) e k) -> return $ In c e (FCase e b1 b2 : k)
  (In CUnit _ k) -> return $ Out VUnit k
  (In (CPair c1 c2) e k) -> return $ In c1 e (FPairR e c2 : k)
  (In (CProj s c) e k) -> return $ In c e (FProj s : k)
  (In (CAbs mem b) e k) -> do
    (xs, body) <- unbind b
    case mem of
      Memo -> do
        cell <- allocateValue (VMap M.empty)
        return $ Out (VClo (Just (cell, [])) e xs body) k
      NoMemo -> return $ Out (VClo Nothing e xs body) k
  (In (CApp c1 c2) e k) -> return $ In c1 e (FArg e c2 : k)
  (In (CType ty) _ k) -> return $ Out (VType ty) k
  (In (CDelay b) e k) -> do
    (xs, cs) <- unbind b
    locs <- allocateRec e (zip (map localName xs) cs)
    return $ Out (foldr (VPair . VRef) VUnit locs) k
  (In (CForce c) e k) -> return $ In c e (FForce : k)
  (In (CTest vars c) e k) -> return $ In c e (FTest (TestVars vars) e : k)
  (Out v (FInj s : k)) -> return $ Out (VInj s v) k
  (Out (VInj L v) (FCase e b1 _ : k)) -> do
    (x, c1) <- unbind b1
    return $ In c1 (Ctx.insert (localName x) v e) k
  (Out (VInj R v) (FCase e _ b2 : k)) -> do
    (x, c2) <- unbind b2
    return $ In c2 (Ctx.insert (localName x) v e) k
  (Out v1 (FPairR e c2 : k)) -> return $ In c2 e (FPairL v1 : k)
  (Out v2 (FPairL v1 : k)) -> return $ Out (VPair v1 v2) k
  (Out (VPair v1 v2) (FProj s : k)) -> return $ Out (selectSide s v1 v2) k
  (Out v (FArg e c2 : k)) -> return $ In c2 e (FApp v : k)
  (Out v (FMemo n sv : k)) -> memoSet n sv v $> Out v k
  (Out v (FApp (VClo mi e [x] b) : k)) -> case mi of
    Nothing -> return $ In b (Ctx.insert (localName x) v e) k
    Just (n, mem) -> do
      let sv = toSimpleValue $ foldr VPair VUnit (v : mem)
      mv <- memoLookup n sv
      case mv of
        Nothing -> return $ In b (Ctx.insert (localName x) v e) (FMemo n sv : k)
        Just v' -> return $ Out v' k
  (Out v (FApp (VClo mi e (x : xs) b) : k)) -> return $ Out (VClo (second (v :) <$> mi) (Ctx.insert (localName x) v e) xs b) k
  (Out v2 (FApp (VConst op) : k)) -> appConst k op v2
  (Out v2 (FApp (VFun f) : k)) -> return $ Out (f v2) k
  -- Annoying to repeat this code, not sure of a better way.
  -- The usual evaluation order (function then argument) doesn't work when
  -- we're applying a test function to randomly generated values.
  (Out (VClo _ e [x] b) (FArgV v : k)) -> return $ In b (Ctx.insert (localName x) v e) k
  (Out (VClo mi e (x : xs) b) (FArgV v : k)) -> return $ Out (VClo mi (Ctx.insert (localName x) v e) xs b) k
  (Out (VConst op) (FArgV v : k)) -> appConst k op v
  (Out (VFun f) (FArgV v : k)) -> return $ Out (f v) k
  (Out (VRef n) (FForce : k)) -> do
    cell <- lkup n
    case cell of
      Nothing -> error $ "impossible: location " ++ show n ++ " not found in memory"
      Just (V v) -> return $ Out v k
      Just (E e t) -> do
        set n Blackhole
        return $ In t e (FUpdate n : k)
      Just Blackhole -> return $ Up InfiniteLoop k
  (Out v (FUpdate n : k)) -> do
    set n (V v)
    return $ Out v k
  (Up err (f@FTest {} : k)) ->
    return $ Out (VProp (VPDone (TestResult False (TestRuntimeError err) emptyTestEnv))) (f : k)
  (Up err (_ : ks)) -> return $ Up err ks
  (Out v (FTest vs e : k)) -> do
    let result = ensureProp v
        res = getTestEnv vs e
    case res of
      Left err -> return $ Up err k
      Right e' -> return $ Out (VProp $ extendPropEnv e' result) k
  _ -> error "Impossible! Bad CESK machine state"

------------------------------------------------------------
-- Interpreting constants
------------------------------------------------------------

arity2 :: (Value -> Value -> a) -> Value -> a
arity2 f (VPair x y) = f x y
arity2 _f _v = error "arity2 on a non-pair!"

arity3 :: (Value -> Value -> Value -> a) -> Value -> a
arity3 f (VPair x (VPair y z)) = f x y z
arity3 _f _v = error "arity3 on a non-triple!"

appConst ::
  Members '[Random, State Mem] r =>
  Cont ->
  Op ->
  Value ->
  Sem r CESK
appConst k = \case
  --------------------------------------------------
  -- Basics

  OCrash -> up . Crash . vlist vchar
  OId -> out
  --------------------------------------------------
  -- Arithmetic

  OAdd -> numOp2 (+) >=> out
  ONeg -> numOp1 negate >=> out
  OSqrt -> numOp1 integerSqrt >=> out
  OFloor -> numOp1 ((% 1) . floor) >=> out
  OCeil -> numOp1 ((% 1) . ceiling) >=> out
  OAbs -> numOp1 abs >=> out
  OMul -> numOp2 (*) >=> out
  ODiv -> numOp2' divOp >>> outWithErr
   where
    divOp :: Member (Error EvalError) r => Rational -> Rational -> Sem r Value
    divOp _ 0 = throw DivByZero
    divOp m n = return $ ratv (m / n)
  OExp -> numOp2' expOp >>> outWithErr
   where
    expOp :: Member (Error EvalError) r => Rational -> Rational -> Sem r Value
    expOp m n
      | m == 0 && n < 0 = throw DivByZero
      | otherwise = return $ ratv (m ^^ numerator n)
  OMod -> numOp2' modOp >>> outWithErr
   where
    modOp :: Member (Error EvalError) r => Rational -> Rational -> Sem r Value
    modOp m n
      | n == 0 = throw DivByZero
      | otherwise = return $ intv (numerator m `mod` numerator n)
  ODivides -> numOp2' (\m n -> return (boolv $ divides m n)) >=> out
  --------------------------------------------------
  -- Number theory

  OIsPrime -> intOp1 (boolv . isPrime) >=> out
  OFactor -> intOp1' primFactor >>> outWithErr
   where
    -- Semantics of the @$factor@ prim: turn a natural number into its
    -- bag of prime factors.  Crash if called on 0, which does not have
    -- a prime factorization.
    primFactor :: Member (Error EvalError) r => Integer -> Sem r Value
    primFactor 0 = throw (Crash "0 has no prime factorization!")
    primFactor n = return . VBag $ map ((intv . unPrime) *** fromIntegral) (factorise n)
  OFrac -> numOp1' (return . primFrac) >=> out
   where
    -- Semantics of the @$frac@ prim: turn a rational number into a pair
    -- of its numerator and denominator.
    primFrac :: Rational -> Value
    primFrac r = VPair (intv (numerator r)) (intv (denominator r))

  --------------------------------------------------
  -- Combinatorics

  OMultinom -> arity2 multinomOp >=> out
   where
    multinomOp :: Value -> Value -> Sem r Value
    multinomOp (vint -> n0) (vlist vint -> ks0) = return . intv $ multinomial n0 ks0
     where
      multinomial :: Integer -> [Integer] -> Integer
      multinomial _ [] = 1
      multinomial n (k' : ks)
        | k' > n = 0
        | otherwise = choose n k' * multinomial (n - k') ks
  OFact -> numOp1' factOp >>> outWithErr
   where
    factOp :: Member (Error EvalError) r => Rational -> Sem r Value
    factOp (numerator -> n)
      | n > fromIntegral (maxBound :: Int) = throw Overflow
      | otherwise = return . intv $ factorial (fromIntegral n)
  OEnum -> out . enumOp
   where
    enumOp :: Value -> Value
    enumOp (VType ty) = listv id (enumerateType ty)
    enumOp v = error $ "Impossible! enumOp on non-type " ++ show v
  OCount -> out . countOp
   where
    countOp :: Value -> Value
    countOp (VType ty) = case countType ty of
      Just num -> VInj R (intv num)
      Nothing -> VNil
    countOp v = error $ "Impossible! countOp on non-type " ++ show v

  --------------------------------------------------
  -- Sequences

  OUntil -> arity2 $ \v1 -> out . ellipsis (Until v1)
  OLookupSeq -> out . oeisLookup
  OExtendSeq -> out . oeisExtend
  --------------------------------------------------
  -- Comparison

  OEq -> arity2 $ \v1 v2 -> out $ boolv (valEq v1 v2)
  OLt -> arity2 $ \v1 v2 -> out $ boolv (valLt v1 v2)
  --------------------------------------------------
  -- Container operations

  OPower -> withBag OPower $ out . VBag . sortNCount . map (first VBag) . choices
   where
    choices :: [(Value, Integer)] -> [([(Value, Integer)], Integer)]
    choices [] = [([], 1)]
    choices ((x, n) : xs) = xs' ++ concatMap (\k' -> map (cons n (x, k')) xs') [1 .. n]
     where
      xs' = choices xs
    cons n (x, k') (zs, m) = ((x, k') : zs, choose n k' * m)
  OBagElem -> arity2 $ \x ->
    withBag OBagElem $
      out . boolv . isJust . find (valEq x) . map fst
  OListElem -> arity2 $ \x -> out . boolv . isJust . find (valEq x) . vlist id
  OEachSet -> arity2 $ \f ->
    withBag OEachSet $
      outWithErr . fmap (VBag . countValues) . mapM (evalApp f . (: []) . fst)
  OEachBag -> arity2 $ \f ->
    withBag OEachBag $
      outWithErr . fmap (VBag . sortNCount) . mapM (\(x, n) -> (,n) <$> evalApp f [x])
  OFilterBag -> arity2 $ \f -> withBag OFilterBag $ \xs ->
    outWithErr $ do
      bs <- mapM (evalApp f . (: []) . fst) xs
      return . VBag . map snd . Prelude.filter (isTrue . fst) $ zip bs xs
   where
    isTrue (VInj R VUnit) = True
    isTrue _ = False
  OMerge -> arity3 $ \f bxs bys ->
    case (bxs, bys) of
      (VBag xs, VBag ys) -> outWithErr (VBag <$> mergeM f xs ys)
      (VBag _, _) -> error $ "Impossible! OMerge on non-VBag " ++ show bys
      _ -> error $ "Impossible! OMerge on non-VBag " ++ show bxs
  OBagUnions -> withBag OBagUnions $ \cts ->
    out . VBag $ sortNCount [(x, m * n) | (VBag xs, n) <- cts, (x, m) <- xs]
  --------------------------------------------------
  -- Container conversions

  OBagToSet -> withBag OBagToSet $ out . VBag . (map . second) (const 1)
  OSetToList -> withBag OSetToList $ out . listv id . map fst
  OBagToList -> withBag OBagToList $ out . listv id . concatMap (uncurry (flip (replicate . fromIntegral)))
  OListToSet -> out . VBag . (map . fmap) (const 1) . countValues . vlist id
  OListToBag -> out . VBag . countValues . vlist id
  OBagToCounts -> withBag OBagToCounts $ out . VBag . map ((,1) . pairv id intv)
  -- Bag (a, N) -> Bag a
  --   Notionally this takes a set of pairs instead of a bag, but operationally we need to
  --   be prepared for a bag, because of the way literal bags desugar, e.g.
  --
  --   Disco> :desugar let x = 3 in ⟅ 'a' # (2 + x), 'b', 'b' ⟆
  --   (λx. bagFromCounts(bag(('a', 2 + x) :: ('b', 1) :: ('b', 1) :: [])))(3)

  OCountsToBag ->
    withBag OCountsToBag $
      out . VBag . sortNCount . map (second (uncurry (*)) . assoc . first (vpair id vint))
   where
    assoc ((a, b), c) = (a, (b, c))
  OUnsafeCountsToBag ->
    withBag OUnsafeCountsToBag $
      out . VBag . map (second (uncurry (*)) . assoc . first (vpair id vint))
   where
    assoc ((a, b), c) = (a, (b, c))

  --------------------------------------------------
  -- Maps

  OMapToSet ->
    withMap OMapToSet $
      out . VBag . map (\(k', v) -> (VPair (fromSimpleValue k') v, 1)) . M.assocs
  OSetToMap ->
    withBag OSetToMap $
      out . VMap . M.fromList . map (convertAssoc . fst)
   where
    convertAssoc (VPair k' v) = (toSimpleValue k', v)
    convertAssoc v = error $ "Impossible! convertAssoc on non-VPair " ++ show v
  OInsert -> arity3 $ \k' v ->
    withMap OInsert $
      out . VMap . M.insert (toSimpleValue k') v
  OLookup -> arity2 $ \k' ->
    withMap OLookup $
      out . toMaybe . M.lookup (toSimpleValue k')
   where
    toMaybe = maybe (VInj L VUnit) (VInj R)
  --------------------------------------------------
  -- Randomness

  ORandom -> arity2 . flip $ \g -> arity2 $ \v1 v2 ->
    let (a, g') = R.randomR (vint v1, vint v2) (vgen g)
     in out $ VPair (intv a) (genv g')
  OSeed -> out . VGen . (R.mkStdGen . fromIntegral . vint)
  --------------------------------------------------
  -- Graph operations

  OVertex -> out . VGraph . Vertex . toSimpleValue
  OOverlay -> arity2 $ withGraph2 OOverlay $ \g1 g2 ->
    out $ VGraph (Overlay g1 g2)
  OConnect -> arity2 $ withGraph2 OConnect $ \g1 g2 ->
    out $ VGraph (Connect g1 g2)
  OSummary -> withGraph OSummary $ out . graphSummary
  --------------------------------------------------
  -- Propositions

  OForall tys -> out . (\v -> VProp (VPSearch SMForall tys v emptyTestEnv))
  OExists tys -> out . (\v -> VProp (VPSearch SMExists tys v emptyTestEnv))
  OHolds -> testProperty Exhaustive >=> resultToBool >>> outWithErr
  ONotProp -> out . VProp . notProp . ensureProp
  OShould op ty -> arity2 $ \v1 v2 ->
    out $ VProp (VPDone (TestResult (valOp op v1 v2) (TestCmp op ty v1 v2) emptyTestEnv))
  OAnd -> arity2 $ \p1 p2 ->
    out $ VProp (VPBin LAnd (ensureProp p1) (ensureProp p2))
  OOr -> arity2 $ \p1 p2 ->
    out $ VProp (VPBin LOr (ensureProp p1) (ensureProp p2))
  OImpl -> arity2 $ \p1 p2 ->
    out $ VProp (VPBin LImpl (ensureProp p1) (ensureProp p2))
  c -> error $ "Unimplemented: appConst " ++ show c
 where
  outWithErr :: Sem (Error EvalError ': r) Value -> Sem r CESK
  outWithErr m = either (`Up` k) (`Out` k) <$> runError m
  out v = return $ Out v k
  up e = return $ Up e k

  withBag :: Op -> ([(Value, Integer)] -> Sem r a) -> Value -> Sem r a
  withBag op f = \case
    VBag xs -> f xs
    v -> error $ "Impossible! " ++ show op ++ " on non-VBag " ++ show v

  withMap :: Op -> (M.Map SimpleValue Value -> Sem r a) -> Value -> Sem r a
  withMap op f = \case
    VMap m -> f m
    v -> error $ "Impossible! " ++ show op ++ " on non-VMap " ++ show v

  withGraph :: Op -> (Graph SimpleValue -> Sem r a) -> Value -> Sem r a
  withGraph op f = \case
    VGraph g -> f g
    v -> error $ "Impossible! " ++ show op ++ " on non-VGraph " ++ show v

  withGraph2 :: Op -> (Graph SimpleValue -> Graph SimpleValue -> Sem r a) -> Value -> Value -> Sem r a
  withGraph2 op f v1 v2 = case (v1, v2) of
    (VGraph g1, VGraph g2) -> f g1 g2
    (_, VGraph _) -> error $ "Impossible! " ++ show op ++ " on non-VGraph " ++ show v1
    _ -> error $ "Impossible! " ++ show op ++ " on non-VGraph " ++ show v2

--------------------------------------------------
-- Arithmetic

intOp1 :: (Integer -> Value) -> Value -> Sem r Value
intOp1 f = intOp1' (return . f)

intOp1' :: (Integer -> Sem r Value) -> Value -> Sem r Value
intOp1' f = f . vint

numOp1 :: (Rational -> Rational) -> Value -> Sem r Value
numOp1 f = numOp1' $ return . ratv . f

numOp1' :: (Rational -> Sem r Value) -> Value -> Sem r Value
numOp1' f (VNum m) = f m
numOp1' _ v = error $ "Impossible! numOp1' on non-VNum " ++ show v

numOp2 :: (Rational -> Rational -> Rational) -> Value -> Sem r Value
numOp2 (#) = numOp2' $ \m n -> return (ratv (m # n))

numOp2' :: (Rational -> Rational -> Sem r Value) -> Value -> Sem r Value
numOp2' (#) =
  arity2 $ \v1 v2 -> case (v1, v2) of
    (VNum n1, VNum n2) -> n1 # n2
    (VNum {}, _) -> error $ "Impossible! numOp2' on non-VNum " ++ show v2
    _ -> error $ "Impossible! numOp2' on non-VNum " ++ show v1

-- | Perform a square root operation. If the program typechecks,
--   then the argument and output will really be Natural.
integerSqrt :: Rational -> Rational
integerSqrt n = integerSqrt' (numerator n) % 1

-- | implementation of `integerSqrt'` taken from the Haskell wiki:
--   https://wiki.haskell.org/Generic_number_type#squareRoot
integerSqrt' :: Integer -> Integer
integerSqrt' 0 = 0
integerSqrt' 1 = 1
integerSqrt' n =
  let twopows = iterate (^! 2) 2
      (lowerRoot, lowerN) =
        last $ takeWhile ((n >=) . snd) $ zip (1 : twopows) twopows
      newtonStep x = div (x + div n x) 2
      iters = InfList.iterate' newtonStep (integerSqrt' (div n lowerN) * lowerRoot)
      isRoot r = r ^! 2 <= n && n < (r + 1) ^! 2
   in InfList.head $ InfList.dropWhile (not . isRoot) iters

-- this operator is used for `integerSqrt'`
(^!) :: Num a => a -> Int -> a
(^!) x n = x ^ n

------------------------------------------------------------
-- Comparison
------------------------------------------------------------

valEq, valLt :: Value -> Value -> Bool
valEq = valOp Eq
valLt = valOp Lt

valOp :: BOp -> Value -> Value -> Bool
valOp op v1 v2 = case op of
  Eq -> valCmp v1 v2 == EQ
  Neq -> valCmp v1 v2 /= EQ
  Lt -> valCmp v1 v2 == LT
  Gt -> valCmp v1 v2 == GT
  Leq -> valCmp v1 v2 /= GT
  Geq -> valCmp v1 v2 /= LT
  Divides -> valDivides v1 v2
  _ -> False

valDivides :: Value -> Value -> Bool
valDivides (VNum r1) (VNum r2) = divides r1 r2
valDivides _ _ = False

divides :: Rational -> Rational -> Bool
divides 0 0 = True
divides 0 _ = False
divides x y = denominator (y / x) == 1

valCmp :: Value -> Value -> Ordering
valCmp (VNum r1) (VNum r2) = compare r1 r2
valCmp (VInj L _) (VInj R _) = LT
valCmp (VInj R _) (VInj L _) = GT
valCmp (VInj L v1) (VInj L v2) = valCmp v1 v2
valCmp (VInj R v1) (VInj R v2) = valCmp v1 v2
valCmp VUnit VUnit = EQ
valCmp (VPair v11 v12) (VPair v21 v22) = valCmp v11 v21 <> valCmp v12 v22
valCmp (VType ty1) (VType ty2) = compare ty1 ty2
valCmp (VBag cs1) (VBag cs2) = compareBags cs1 cs2
valCmp (VMap m1) (VMap m2) = compareMaps (M.assocs m1) (M.assocs m2)
valCmp (VGraph g1) (VGraph g2) = valCmp (graphSummary g1) (graphSummary g2)
valCmp v1 v2 = error $ "valCmp\n  " ++ take 100 (show v1) ++ "...\n  " ++ take 100 (show v2) ++ "..."

compareBags :: [(Value, Integer)] -> [(Value, Integer)] -> Ordering
compareBags [] [] = EQ
compareBags [] _ = LT
compareBags _ [] = GT
compareBags ((x, xn) : xs) ((y, yn) : ys) =
  valCmp x y <> compare xn yn <> compareBags xs ys

compareMaps :: [(SimpleValue, Value)] -> [(SimpleValue, Value)] -> Ordering
compareMaps [] [] = EQ
compareMaps [] _ = LT
compareMaps _ [] = GT
compareMaps ((k1, v1) : as1) ((k2, v2) : as2) =
  valCmp (fromSimpleValue k1) (fromSimpleValue k2) <> valCmp v1 v2 <> compareMaps as1 as2

------------------------------------------------------------
-- Polynomial sequences [a,b,c,d .. e]
------------------------------------------------------------

ellipsis :: Ellipsis Value -> Value -> Value
ellipsis (fmap vrat -> end) (vlist vrat -> rs) = listv ratv $ enumEllipsis rs end

enumEllipsis :: (Enum a, Num a, Ord a) => [a] -> Ellipsis a -> [a]
enumEllipsis [] _ = error "Impossible! Disco.Interpret.CESK.enumEllipsis []"
enumEllipsis [x] (Until y)
  | x <= y = [x .. y]
  | otherwise = [x, pred x .. y]
enumEllipsis xs (Until y)
  | d > 0 = takeWhile (<= y) nums
  | d < 0 = takeWhile (>= y) nums
  | otherwise = nums
 where
  d = constdiff xs
  nums = babbage xs

-- | Extend a sequence infinitely by interpolating it as a polynomial
--   sequence, via forward differences.  Essentially the same
--   algorithm used by Babbage's famous Difference Engine.
babbage :: Num a => [a] -> [a]
babbage [] = []
babbage [x] = repeat x
babbage (x : xs) = scanl (+) x (babbage (diff (x : xs)))

-- | Compute the forward difference of the given sequence, that is,
--   differences of consecutive pairs of elements.
diff :: Num a => [a] -> [a]
diff xs = zipWith (-) (drop 1 xs) xs

-- | Take forward differences until the result is constant, and return
--   the constant.  The sign of the constant difference tells us the
--   limiting behavior of the sequence.
constdiff :: (Eq a, Num a) => [a] -> a
constdiff [] = error "Impossible! Disco.Interpret.Core.constdiff []"
constdiff (x : xs)
  | all (== x) xs = x
  | otherwise = constdiff (diff (x : xs))

------------------------------------------------------------
-- OEIS
------------------------------------------------------------

-- | Looks up a sequence of integers in OEIS.
--   Returns 'left()' if the sequence is unknown in OEIS,
--   otherwise 'right "https://oeis.org/<oeis_sequence_id>"'
oeisLookup :: Value -> Value
oeisLookup (vlist vint -> ns) = maybe VNil parseResult (lookupSeq (SubSeq ns))
 where
  parseResult r = VInj R (listv charv ("https://oeis.org/" ++ T.unpack (number r)))

-- | Extends a Disco integer list with data from a known OEIS
--   sequence.  Returns a list of integers upon success, otherwise the
--   original list (unmodified).
oeisExtend :: Value -> Value
oeisExtend = listv intv . extendSeq . vlist vint

------------------------------------------------------------
-- Normalizing bags/sets
------------------------------------------------------------

-- | Given a list of disco values, sort and collate them into a list
--   pairing each unique value with its count.  Used to
--   construct/normalize bags and sets.  Prerequisite: the values must
--   be comparable.
countValues :: [Value] -> [(Value, Integer)]
countValues = sortNCount . map (,1)

-- | Normalize a list of values where each value is paired with a
--   count, but there could be duplicate values.  This function uses
--   merge sort to sort the values, adding the counts of multiple
--   instances of the same value.  Prerequisite: the values must be
--   comparable.
sortNCount :: [(Value, Integer)] -> [(Value, Integer)]
sortNCount [] = []
sortNCount [x] = [x]
sortNCount xs = merge (+) (sortNCount firstHalf) (sortNCount secondHalf)
 where
  (firstHalf, secondHalf) = splitAt (length xs `div` 2) xs

-- | Generic function for merging two sorted, count-annotated lists of
--   type @[(a,Integer)]@ a la merge sort, using the given comparison
--   function, and using the provided count combining function to
--   decide what count to assign to each element of the output.  For
--   example, @(+)@ corresponds to bag union; @min@ corresponds to
--   intersection; and so on.
merge ::
  (Integer -> Integer -> Integer) ->
  [(Value, Integer)] ->
  [(Value, Integer)] ->
  [(Value, Integer)]
merge g = go
 where
  go [] [] = []
  go [] ((y, n) : ys) = mergeCons y 0 n (go [] ys)
  go ((x, n) : xs) [] = mergeCons x n 0 (go xs [])
  go ((x, n1) : xs) ((y, n2) : ys) = case valCmp x y of
    LT -> mergeCons x n1 0 (go xs ((y, n2) : ys))
    EQ -> mergeCons x n1 n2 (go xs ys)
    GT -> mergeCons y 0 n2 (go ((x, n1) : xs) ys)

  mergeCons a m1 m2 zs = case g m1 m2 of
    0 -> zs
    n -> (a, n) : zs

mergeM ::
  Members '[Random, Error EvalError, State Mem] r =>
  Value ->
  [(Value, Integer)] ->
  [(Value, Integer)] ->
  Sem r [(Value, Integer)]
mergeM g = go
 where
  go [] [] = return []
  go [] ((y, n) : ys) = mergeCons y 0 n =<< go [] ys
  go ((x, n) : xs) [] = mergeCons x n 0 =<< go xs []
  go ((x, n1) : xs) ((y, n2) : ys) = case valCmp x y of
    LT -> mergeCons x n1 0 =<< go xs ((y, n2) : ys)
    EQ -> mergeCons x n1 n2 =<< go xs ys
    GT -> mergeCons y 0 n2 =<< go ((x, n1) : xs) ys

  mergeCons a m1 m2 zs = do
    nm <- evalApp g [VPair (intv m1) (intv m2)]
    return $ case nm of
      VNum 0 -> zs
      VNum n -> (a, numerator n) : zs
      v -> error $ "Impossible! merge function in mergeM returned non-VNum " ++ show v

------------------------------------------------------------
-- Graphs
------------------------------------------------------------

graphSummary :: Graph SimpleValue -> Value
graphSummary = toDiscoAdjMap . reifyGraph
 where
  reifyGraph :: Graph SimpleValue -> [(SimpleValue, [SimpleValue])]
  reifyGraph =
    AdjMap.adjacencyList . foldg AdjMap.empty AdjMap.vertex AdjMap.overlay AdjMap.connect

  toDiscoAdjMap :: [(SimpleValue, [SimpleValue])] -> Value
  toDiscoAdjMap =
    VMap . M.fromList . map (second (VBag . countValues . map fromSimpleValue))

------------------------------------------------------------
-- Propositions / tests
------------------------------------------------------------

resultToBool :: Member (Error EvalError) r => TestResult -> Sem r Value
resultToBool (TestResult _ (TestRuntimeError e) _) = throw e
resultToBool (TestResult b _ _) = return $ boolv b

notProp :: ValProp -> ValProp
notProp (VPDone r) = VPDone (invertPropResult r)
notProp (VPSearch sm tys p e) = VPSearch (invertMotive sm) tys p e
notProp (VPBin LAnd vp1 vp2) = VPBin LOr (notProp vp1) (notProp vp2)
notProp (VPBin LOr vp1 vp2) = VPBin LAnd (notProp vp1) (notProp vp2)
notProp (VPBin LImpl vp1 vp2) = VPBin LAnd vp1 (notProp vp2)

-- | Convert a @Value@ to a @ValProp@, embedding booleans if necessary.
ensureProp :: Value -> ValProp
ensureProp (VProp p) = p
ensureProp (VInj L _) = VPDone (TestResult False TestBool emptyTestEnv)
ensureProp (VInj R _) = VPDone (TestResult True TestBool emptyTestEnv)
ensureProp _ = error "ensureProp: non-prop value"

combineTestResults :: LOp -> TestResult -> TestResult -> TestResult
combineTestResults op tr1@(TestResult b1 _ e1) tr2@(TestResult b2 _ e2) =
  TestResult (interpLOp op b1 b2) (TestBin op tr1 tr2) (mergeTestEnv e1 e2)

testProperty ::
  Members '[Random, State Mem] r =>
  SearchType ->
  Value ->
  Sem r TestResult
testProperty initialSt = checkProp . ensureProp
 where
  checkProp ::
    Members '[Random, State Mem] r =>
    ValProp ->
    Sem r TestResult
  checkProp (VPDone r) = return r
  checkProp (VPBin op vp1 vp2) = do
    tr1 <- checkProp vp1
    tr2 <- checkProp vp2
    return $ combineTestResults op tr1 tr2
  checkProp (VPSearch sm tys f e) =
    extendResultEnv e <$> (generateSamples initialSt vals >>= go)
   where
    vals = enumTypes tys
    (SearchMotive (whenFound, wantsSuccess)) = sm

    go ::
      Members '[Random, State Mem] r =>
      ([[Value]], SearchType) ->
      Sem r TestResult
    go ([], st) = return $ TestResult (not whenFound) (TestNotFound st) emptyTestEnv
    go (x : xs, st) = do
      mprop <- runError (ensureProp <$> evalApp f x)
      case mprop of
        Left err -> return $ TestResult False (TestRuntimeError err) emptyTestEnv
        Right (VPDone r) -> continue st xs r
        Right prop -> checkProp prop >>= continue st xs

    continue ::
      Members '[Random, State Mem] r =>
      SearchType ->
      [[Value]] ->
      TestResult ->
      Sem r TestResult
    continue st xs r@(TestResult _ _ e')
      | testIsError r = return r
      | testIsOk r == wantsSuccess =
          return $ TestResult whenFound (TestFound r) e'
      | otherwise = go (xs, st)

evalApp ::
  Members '[Random, Error EvalError, State Mem] r =>
  Value ->
  [Value] ->
  Sem r Value
evalApp f xs =
  runFresh (runCESK (Out f (map FArgV xs))) >>= either throw return

runTest ::
  Members '[Random, Error EvalError, Input Env, State Mem] r =>
  Int ->
  AProperty ->
  Sem r TestResult
runTest n p = testProperty (Randomized n' n') =<< eval (compileProperty p)
 where
  n' = fromIntegral (n `div` 2)

------------------------------------------------------------
-- Top-level evaluation
------------------------------------------------------------

eval :: Members '[Random, Error EvalError, Input Env, State Mem] r => Core -> Sem r Value
eval c = do
  e <- input @Env
  runFresh (runCESK (In c e [])) >>= either throw return