AlgReal.hs

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveLift #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      :   Grisette.Internal.SymPrim.AlgReal
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.SymPrim.AlgReal
  ( AlgRealPoly (..),
    UnsupportedAlgRealOperation (..),
    toSBVAlgReal,
    fromSBVAlgReal,
    RealPoint (..),
    AlgReal (..),
  )
where

import Control.DeepSeq (NFData)
import Control.Exception (Exception, throw)
import qualified Data.Binary as Binary
import Data.Bytes.Serial (Serial (deserialize, serialize))
import Data.Hashable (Hashable)
import qualified Data.SBV as SBV
import qualified Data.SBV.Internals as SBV
import qualified Data.Serialize as Cereal
import GHC.Generics (Generic)
import Grisette.Internal.Core.Data.Class.Function (Apply (FunType, apply))
import Language.Haskell.TH.Syntax (Lift)
import Test.QuickCheck (Arbitrary)
import Test.QuickCheck.Arbitrary (Arbitrary (arbitrary))

-- | A univariate polynomial with integer coefficients.
--
-- For instance, @5x^3+2x-5@ is represented as
-- @v'AlgRealPoly' [(5, 3), (2, 1), (-5, 0)]@.
newtype AlgRealPoly = AlgRealPoly [(Integer, Integer)]
  deriving (Eq, Generic, Lift)
  deriving newtype (Hashable, NFData)
  deriving anyclass (Serial)

instance Cereal.Serialize AlgRealPoly where
  put = serialize
  get = deserialize

instance Binary.Binary AlgRealPoly where
  put = serialize
  get = deserialize

-- | Boundary point for real intervals.
data RealPoint
  = -- | Open point.
    OpenPoint Rational
  | -- | Closed point.
    ClosedPoint Rational
  deriving (Eq, Generic, Lift)
  deriving anyclass (Hashable, NFData, Serial)

instance Cereal.Serialize RealPoint where
  put = serialize
  get = deserialize

instance Binary.Binary RealPoint where
  put = serialize
  get = deserialize

toSBVRealPoint :: RealPoint -> SBV.RealPoint Rational
toSBVRealPoint (OpenPoint r) = SBV.OpenPoint r
toSBVRealPoint (ClosedPoint r) = SBV.ClosedPoint r

fromSBVRealPoint :: SBV.RealPoint Rational -> RealPoint
fromSBVRealPoint (SBV.OpenPoint r) = OpenPoint r
fromSBVRealPoint (SBV.ClosedPoint r) = ClosedPoint r

-- | Algebraic real numbers. The representation can be abstract for
-- roots-of-polynomials or intervals.
data AlgReal where
  -- | Exact rational number.
  AlgExactRational :: Rational -> AlgReal
  -- | Inexact rational numbers. SMT-solver return it with ? at the end.
  AlgInexactRational :: Rational -> AlgReal
  -- | Algebraic real number as a root of a polynomial.
  AlgPolyRoot ::
    -- | Which root is it?
    Integer ->
    -- | Polynomial defining equation.
    AlgRealPoly ->
    -- | Approximate decimal representation.
    Maybe String ->
    AlgReal
  -- | Interval with low and high bounds.
  AlgInterval ::
    -- | Lower bound.
    RealPoint ->
    -- | Upper bound.
    RealPoint ->
    AlgReal
  deriving (Generic, Lift)
  deriving anyclass (Hashable, NFData, Serial)

instance Cereal.Serialize AlgReal where
  put = serialize
  get = deserialize

instance Binary.Binary AlgReal where
  put = serialize
  get = deserialize

-- | Convert algebraic real numbers to SBV's algebraic real numbers.
toSBVAlgReal :: AlgReal -> SBV.AlgReal
toSBVAlgReal (AlgExactRational r) = SBV.AlgRational True r
toSBVAlgReal (AlgInexactRational r) = SBV.AlgRational False r
toSBVAlgReal (AlgPolyRoot i (AlgRealPoly ps) approx) =
  SBV.AlgPolyRoot (i, SBV.AlgRealPoly ps) approx
toSBVAlgReal (AlgInterval l u) =
  SBV.AlgInterval (toSBVRealPoint l) (toSBVRealPoint u)

-- | Convert SBV's algebraic real numbers to algebraic real numbers.
fromSBVAlgReal :: SBV.AlgReal -> AlgReal
fromSBVAlgReal (SBV.AlgRational True r) = AlgExactRational r
fromSBVAlgReal (SBV.AlgRational False r) = AlgInexactRational r
fromSBVAlgReal (SBV.AlgPolyRoot (i, SBV.AlgRealPoly ps) approx) =
  AlgPolyRoot i (AlgRealPoly ps) approx
fromSBVAlgReal (SBV.AlgInterval l u) =
  AlgInterval (fromSBVRealPoint l) (fromSBVRealPoint u)

instance Show AlgReal where
  show r = show $ toSBVAlgReal r

-- | Exception for unsupported operations on algebraic real numbers.
--
-- We only support operations on exact rationals.
data UnsupportedAlgRealOperation = UnsupportedAlgRealOperation
  { op :: String,
    msg :: String
  }
  deriving anyclass (Exception)

instance Show UnsupportedAlgRealOperation where
  show (UnsupportedAlgRealOperation op msg) =
    "AlgReal."
      ++ op
      ++ ": unsupported operation on algebraic rationals, only support exact "
      ++ "rationals"
      ++ ": "
      ++ msg

op1 :: String -> (Rational -> Rational) -> AlgReal -> AlgReal
op1 _ f (AlgExactRational r) = AlgExactRational $ f r
op1 name _ r =
  throw
    UnsupportedAlgRealOperation {op = name, msg = show r}

op2 ::
  String ->
  (Rational -> Rational -> Rational) ->
  AlgReal ->
  AlgReal ->
  AlgReal
op2 _ f (AlgExactRational l) (AlgExactRational r) = AlgExactRational $ f l r
op2 name _ l r =
  throw
    UnsupportedAlgRealOperation
      { op = name,
        msg = show l <> " and " <> show r
      }

instance Eq AlgReal where
  (AlgExactRational l) == (AlgExactRational r) = l == r
  l == r =
    throw $
      UnsupportedAlgRealOperation "==" $
        show l <> " and " <> show r

instance Ord AlgReal where
  compare (AlgExactRational l) (AlgExactRational r) = compare l r
  compare l r =
    throw $
      UnsupportedAlgRealOperation "compare" $
        show l <> " and " <> show r

instance Num AlgReal where
  (+) = op2 "+" (+)
  (*) = op2 "*" (*)
  (-) = op2 "-" (-)
  negate = op1 "negate" negate
  abs = op1 "abs" abs
  signum = op1 "signum" signum
  fromInteger = AlgExactRational . fromInteger

-- | Unlike sbv, we throw the error when divided by zero happens
instance Fractional AlgReal where
  (/) = op2 "/" (/)
  fromRational = AlgExactRational

instance Real AlgReal where
  toRational (AlgExactRational r) = r
  toRational r =
    throw $
      UnsupportedAlgRealOperation "toRational" $
        show r

instance Arbitrary AlgReal where
  arbitrary = AlgExactRational <$> arbitrary

instance Apply AlgReal where
  type FunType AlgReal = AlgReal
  apply = id