{-|
Module      : What4.Protocol.PolyRoot
Description : Representation for algebraic reals
Copyright   : (c) Galois Inc, 2016-2020
License     : BSD3
Maintainer  : jhendrix@galois.com

Defines a numeric data-type where each number is represented as the root of a
polynomial over a single variable.

This currently only defines operations for parsing the roots from the format
generated by Yices, and evaluating a polynomial over rational coefficients
to the rational derived from the closest double.
-}

{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
module What4.Protocol.PolyRoot
  ( Root
  , approximate
  , fromYicesText
  , parseYicesRoot
  ) where

import           Control.Applicative
import           Control.Lens
import qualified Data.Attoparsec.Text as Atto
import qualified Data.Map as Map
import           Data.Ratio
import           Data.Text (Text)
import qualified Data.Text as Text

import qualified Data.Vector as V
import           Prettyprinter as PP

atto_angle :: Atto.Parser a -> Atto.Parser a
atto_angle :: Parser a -> Parser a
atto_angle Parser a
p = Char -> Parser Char
Atto.char Char
'<' Parser Char -> Parser a -> Parser a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser a
p Parser a -> Parser Char -> Parser a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* Char -> Parser Char
Atto.char Char
'>'

atto_paren :: Atto.Parser a -> Atto.Parser a
atto_paren :: Parser a -> Parser a
atto_paren Parser a
p = Char -> Parser Char
Atto.char Char
'(' Parser Char -> Parser a -> Parser a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser a
p Parser a -> Parser Char -> Parser a
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f a
<* Char -> Parser Char
Atto.char Char
')'

-- | A polynomial with one variable.
newtype SingPoly coef = SingPoly (V.Vector coef)
  deriving (a -> SingPoly b -> SingPoly a
(a -> b) -> SingPoly a -> SingPoly b
(forall a b. (a -> b) -> SingPoly a -> SingPoly b)
-> (forall a b. a -> SingPoly b -> SingPoly a) -> Functor SingPoly
forall a b. a -> SingPoly b -> SingPoly a
forall a b. (a -> b) -> SingPoly a -> SingPoly b
forall (f :: Type -> Type).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> SingPoly b -> SingPoly a
$c<$ :: forall a b. a -> SingPoly b -> SingPoly a
fmap :: (a -> b) -> SingPoly a -> SingPoly b
$cfmap :: forall a b. (a -> b) -> SingPoly a -> SingPoly b
Functor, a -> SingPoly a -> Bool
SingPoly m -> m
SingPoly a -> [a]
SingPoly a -> Bool
SingPoly a -> Int
SingPoly a -> a
SingPoly a -> a
SingPoly a -> a
SingPoly a -> a
(a -> m) -> SingPoly a -> m
(a -> m) -> SingPoly a -> m
(a -> b -> b) -> b -> SingPoly a -> b
(a -> b -> b) -> b -> SingPoly a -> b
(b -> a -> b) -> b -> SingPoly a -> b
(b -> a -> b) -> b -> SingPoly a -> b
(a -> a -> a) -> SingPoly a -> a
(a -> a -> a) -> SingPoly a -> a
(forall m. Monoid m => SingPoly m -> m)
-> (forall m a. Monoid m => (a -> m) -> SingPoly a -> m)
-> (forall m a. Monoid m => (a -> m) -> SingPoly a -> m)
-> (forall a b. (a -> b -> b) -> b -> SingPoly a -> b)
-> (forall a b. (a -> b -> b) -> b -> SingPoly a -> b)
-> (forall b a. (b -> a -> b) -> b -> SingPoly a -> b)
-> (forall b a. (b -> a -> b) -> b -> SingPoly a -> b)
-> (forall a. (a -> a -> a) -> SingPoly a -> a)
-> (forall a. (a -> a -> a) -> SingPoly a -> a)
-> (forall a. SingPoly a -> [a])
-> (forall a. SingPoly a -> Bool)
-> (forall a. SingPoly a -> Int)
-> (forall a. Eq a => a -> SingPoly a -> Bool)
-> (forall a. Ord a => SingPoly a -> a)
-> (forall a. Ord a => SingPoly a -> a)
-> (forall a. Num a => SingPoly a -> a)
-> (forall a. Num a => SingPoly a -> a)
-> Foldable SingPoly
forall a. Eq a => a -> SingPoly a -> Bool
forall a. Num a => SingPoly a -> a
forall a. Ord a => SingPoly a -> a
forall m. Monoid m => SingPoly m -> m
forall a. SingPoly a -> Bool
forall a. SingPoly a -> Int
forall a. SingPoly a -> [a]
forall a. (a -> a -> a) -> SingPoly a -> a
forall m a. Monoid m => (a -> m) -> SingPoly a -> m
forall b a. (b -> a -> b) -> b -> SingPoly a -> b
forall a b. (a -> b -> b) -> b -> SingPoly a -> b
forall (t :: Type -> Type).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: SingPoly a -> a
$cproduct :: forall a. Num a => SingPoly a -> a
sum :: SingPoly a -> a
$csum :: forall a. Num a => SingPoly a -> a
minimum :: SingPoly a -> a
$cminimum :: forall a. Ord a => SingPoly a -> a
maximum :: SingPoly a -> a
$cmaximum :: forall a. Ord a => SingPoly a -> a
elem :: a -> SingPoly a -> Bool
$celem :: forall a. Eq a => a -> SingPoly a -> Bool
length :: SingPoly a -> Int
$clength :: forall a. SingPoly a -> Int
null :: SingPoly a -> Bool
$cnull :: forall a. SingPoly a -> Bool
toList :: SingPoly a -> [a]
$ctoList :: forall a. SingPoly a -> [a]
foldl1 :: (a -> a -> a) -> SingPoly a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> SingPoly a -> a
foldr1 :: (a -> a -> a) -> SingPoly a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> SingPoly a -> a
foldl' :: (b -> a -> b) -> b -> SingPoly a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> SingPoly a -> b
foldl :: (b -> a -> b) -> b -> SingPoly a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> SingPoly a -> b
foldr' :: (a -> b -> b) -> b -> SingPoly a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> SingPoly a -> b
foldr :: (a -> b -> b) -> b -> SingPoly a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> SingPoly a -> b
foldMap' :: (a -> m) -> SingPoly a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> SingPoly a -> m
foldMap :: (a -> m) -> SingPoly a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> SingPoly a -> m
fold :: SingPoly m -> m
$cfold :: forall m. Monoid m => SingPoly m -> m
Foldable, Functor SingPoly
Foldable SingPoly
Functor SingPoly
-> Foldable SingPoly
-> (forall (f :: Type -> Type) a b.
    Applicative f =>
    (a -> f b) -> SingPoly a -> f (SingPoly b))
-> (forall (f :: Type -> Type) a.
    Applicative f =>
    SingPoly (f a) -> f (SingPoly a))
-> (forall (m :: Type -> Type) a b.
    Monad m =>
    (a -> m b) -> SingPoly a -> m (SingPoly b))
-> (forall (m :: Type -> Type) a.
    Monad m =>
    SingPoly (m a) -> m (SingPoly a))
-> Traversable SingPoly
(a -> f b) -> SingPoly a -> f (SingPoly b)
forall (t :: Type -> Type).
Functor t
-> Foldable t
-> (forall (f :: Type -> Type) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: Type -> Type) a.
    Applicative f =>
    t (f a) -> f (t a))
-> (forall (m :: Type -> Type) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: Type -> Type) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: Type -> Type) a.
Monad m =>
SingPoly (m a) -> m (SingPoly a)
forall (f :: Type -> Type) a.
Applicative f =>
SingPoly (f a) -> f (SingPoly a)
forall (m :: Type -> Type) a b.
Monad m =>
(a -> m b) -> SingPoly a -> m (SingPoly b)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> SingPoly a -> f (SingPoly b)
sequence :: SingPoly (m a) -> m (SingPoly a)
$csequence :: forall (m :: Type -> Type) a.
Monad m =>
SingPoly (m a) -> m (SingPoly a)
mapM :: (a -> m b) -> SingPoly a -> m (SingPoly b)
$cmapM :: forall (m :: Type -> Type) a b.
Monad m =>
(a -> m b) -> SingPoly a -> m (SingPoly b)
sequenceA :: SingPoly (f a) -> f (SingPoly a)
$csequenceA :: forall (f :: Type -> Type) a.
Applicative f =>
SingPoly (f a) -> f (SingPoly a)
traverse :: (a -> f b) -> SingPoly a -> f (SingPoly b)
$ctraverse :: forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> SingPoly a -> f (SingPoly b)
$cp2Traversable :: Foldable SingPoly
$cp1Traversable :: Functor SingPoly
Traversable, Int -> SingPoly coef -> ShowS
[SingPoly coef] -> ShowS
SingPoly coef -> String
(Int -> SingPoly coef -> ShowS)
-> (SingPoly coef -> String)
-> ([SingPoly coef] -> ShowS)
-> Show (SingPoly coef)
forall coef. Show coef => Int -> SingPoly coef -> ShowS
forall coef. Show coef => [SingPoly coef] -> ShowS
forall coef. Show coef => SingPoly coef -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SingPoly coef] -> ShowS
$cshowList :: forall coef. Show coef => [SingPoly coef] -> ShowS
show :: SingPoly coef -> String
$cshow :: forall coef. Show coef => SingPoly coef -> String
showsPrec :: Int -> SingPoly coef -> ShowS
$cshowsPrec :: forall coef. Show coef => Int -> SingPoly coef -> ShowS
Show)

instance (Ord coef, Num coef, Pretty coef) => Pretty (SingPoly coef) where
  pretty :: SingPoly coef -> Doc ann
pretty (SingPoly Vector coef
v) =
    case (coef -> Bool) -> Vector coef -> Maybe Int
forall a. (a -> Bool) -> Vector a -> Maybe Int
V.findIndex (coef -> coef -> Bool
forall a. Eq a => a -> a -> Bool
/= coef
0) Vector coef
v of
      Maybe Int
Nothing -> String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"0"
      Just Int
j -> Int -> Doc ann
forall ann. Int -> Doc ann
go (Vector coef -> Int
forall a. Vector a -> Int
V.length Vector coef
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        where ppc :: a -> Doc ann
ppc a
c | a
c a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
c)
                    | Bool
otherwise = a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
c

              ppi :: a -> Doc ann
ppi a
1 = String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"*x"
              ppi a
i = String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"*x^" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
i

              go :: Int -> Doc ann
go Int
0 = coef -> Doc ann
forall a ann. (Ord a, Num a, Pretty a) => a -> Doc ann
ppc (Vector coef
v Vector coef -> Int -> coef
forall a. Vector a -> Int -> a
V.! Int
0)
              go Int
i | Int -> Bool -> Bool
seq Int
i Bool
False = String -> Doc ann
forall a. HasCallStack => String -> a
error String
"pretty SingPoly"
                   | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = coef -> Doc ann
forall a ann. (Ord a, Num a, Pretty a) => a -> Doc ann
ppc (Vector coef
v Vector coef -> Int -> coef
forall a. Vector a -> Int -> a
V.! Int
i) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann
forall a ann. (Eq a, Num a, Pretty a) => a -> Doc ann
ppi Int
i
                   | Vector coef
v Vector coef -> Int -> coef
forall a. Vector a -> Int -> a
V.! Int
i coef -> coef -> Bool
forall a. Eq a => a -> a -> Bool
== coef
0 = Int -> Doc ann
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                   | Bool
otherwise = coef -> Doc ann
forall a ann. (Ord a, Num a, Pretty a) => a -> Doc ann
ppc (Vector coef
v Vector coef -> Int -> coef
forall a. Vector a -> Int -> a
V.! Int
i) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann
forall a ann. (Eq a, Num a, Pretty a) => a -> Doc ann
ppi Int
i Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
"+" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Int -> Doc ann
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

fromList :: [c] -> SingPoly c
fromList :: [c] -> SingPoly c
fromList = Vector c -> SingPoly c
forall coef. Vector coef -> SingPoly coef
SingPoly (Vector c -> SingPoly c) -> ([c] -> Vector c) -> [c] -> SingPoly c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [c] -> Vector c
forall a. [a] -> Vector a
V.fromList

-- | Create a polyomial from a map from powers to coefficient.
fromMap :: (Eq c, Num c) => Map.Map Int c -> SingPoly c
fromMap :: Map Int c -> SingPoly c
fromMap Map Int c
m0 = Vector c -> SingPoly c
forall coef. Vector coef -> SingPoly coef
SingPoly (Int -> (Int -> c) -> Vector c
forall a. Int -> (Int -> a) -> Vector a
V.generate (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> c
f)
  where m :: Map Int c
m = (c -> Bool) -> Map Int c -> Map Int c
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (c -> c -> Bool
forall a. Eq a => a -> a -> Bool
/= c
0) Map Int c
m0
        (Int
n,c
_) = Map Int c -> (Int, c)
forall k a. Map k a -> (k, a)
Map.findMax Map Int c
m
        f :: Int -> c
f Int
i   = c -> Int -> Map Int c -> c
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault c
0 Int
i Map Int c
m

-- | Parse a positive monomial
pos_mono :: Integral c => Atto.Parser (c, Int)
pos_mono :: Parser (c, Int)
pos_mono = (,) (c -> Int -> (c, Int))
-> Parser Text c -> Parser Text (Int -> (c, Int))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Text c
forall a. Integral a => Parser a
Atto.decimal Parser Text (Int -> (c, Int)) -> Parser Text Int -> Parser (c, Int)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Parser Text Int
times_x
  where times_x :: Atto.Parser Int
        times_x :: Parser Text Int
times_x = (Char -> Parser Char
Atto.char Char
'*' Parser Char -> Parser Char -> Parser Char
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Char -> Parser Char
Atto.char Char
'x' Parser Char -> Parser Text Int -> Parser Text Int
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser Text Int
expon) Parser Text Int -> Parser Text Int -> Parser Text Int
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> Int -> Parser Text Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
0

        -- Parse explicit exponent or return 1
        expon :: Atto.Parser Int
        expon :: Parser Text Int
expon = (Char -> Parser Char
Atto.char Char
'^' Parser Char -> Parser Text Int -> Parser Text Int
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser Text Int
forall a. Integral a => Parser a
Atto.decimal) Parser Text Int -> Parser Text Int -> Parser Text Int
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> Int -> Parser Text Int
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Int
1


-- | Parses a monomial and returns the coefficient and power
mono :: Integral c => Atto.Parser (c, Int)
mono :: Parser (c, Int)
mono = Parser (c, Int) -> Parser (c, Int)
forall a. Parser a -> Parser a
atto_paren (Char -> Parser Char
Atto.char Char
'-' Parser Char -> Parser (c, Int) -> Parser (c, Int)
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> (ASetter (c, Int) (c, Int) c c -> (c -> c) -> (c, Int) -> (c, Int)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (c, Int) (c, Int) c c
forall s t a b. Field1 s t a b => Lens s t a b
_1 c -> c
forall a. Num a => a -> a
negate ((c, Int) -> (c, Int)) -> Parser (c, Int) -> Parser (c, Int)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser (c, Int)
forall c. Integral c => Parser (c, Int)
pos_mono))
     Parser (c, Int) -> Parser (c, Int) -> Parser (c, Int)
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> Parser (c, Int)
forall c. Integral c => Parser (c, Int)
pos_mono

parseYicesPoly :: Integral c => Atto.Parser (SingPoly c)
parseYicesPoly :: Parser (SingPoly c)
parseYicesPoly = do
     (c
c,Int
p) <- Parser (c, Int)
forall c. Integral c => Parser (c, Int)
mono
     Map Int c -> Parser (SingPoly c)
forall a. Integral a => Map Int a -> Parser Text (SingPoly a)
go (Int -> c -> Map Int c
forall k a. k -> a -> Map k a
Map.singleton Int
p c
c)
  where go :: Map Int a -> Parser Text (SingPoly a)
go Map Int a
m = Map Int a -> Parser Text (SingPoly a)
next Map Int a
m Parser Text (SingPoly a)
-> Parser Text (SingPoly a) -> Parser Text (SingPoly a)
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> SingPoly a -> Parser Text (SingPoly a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Map Int a -> SingPoly a
forall c. (Eq c, Num c) => Map Int c -> SingPoly c
fromMap Map Int a
m)
        next :: Map Int a -> Parser Text (SingPoly a)
next Map Int a
m = Map Int a -> Parser Text (SingPoly a) -> Parser Text (SingPoly a)
seq Map Int a
m (Parser Text (SingPoly a) -> Parser Text (SingPoly a))
-> Parser Text (SingPoly a) -> Parser Text (SingPoly a)
forall a b. (a -> b) -> a -> b
$ do
          Char
_ <- Char -> Parser Char
Atto.char Char
' ' Parser Char -> Parser Char -> Parser Char
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Char -> Parser Char
Atto.char Char
'+' Parser Char -> Parser Char -> Parser Char
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Char -> Parser Char
Atto.char Char
' '
          (a
c,Int
p) <- Parser (a, Int)
forall c. Integral c => Parser (c, Int)
mono
          Map Int a -> Parser Text (SingPoly a)
go ((a -> a -> a) -> Int -> a -> Map Int a -> Map Int a
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Int
p a
c Map Int a
m)


-- | Evaluate polynomial at a specific value.
--
-- Note that due to rounding, the result may not be exact when using
-- finite precision arithmetic.
eval :: forall c . Num c => SingPoly c -> c -> c
eval :: SingPoly c -> c -> c
eval (SingPoly Vector c
v) c
c = Int -> c -> c -> c
f Int
0 c
1 c
0
  where -- f takes an index, the current power, and the current sum.
        f :: Int -> c -> c -> c
        f :: Int -> c -> c -> c
f Int
i c
p c
s
          | c -> Bool -> Bool
seq c
p (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ c -> Bool -> Bool
seq c
s (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Bool
False = String -> c
forall a. HasCallStack => String -> a
error String
"internal error: Poly.eval"
          | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector c -> Int
forall a. Vector a -> Int
V.length Vector c
v = Int -> c -> c -> c
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (c
p c -> c -> c
forall a. Num a => a -> a -> a
* c
c) (c
s c -> c -> c
forall a. Num a => a -> a -> a
+ c
p c -> c -> c
forall a. Num a => a -> a -> a
* (Vector c
v Vector c -> Int -> c
forall a. Vector a -> Int -> a
V.! Int
i))
          | Bool
otherwise = c
s

data Root c = Root { Root c -> SingPoly c
rootPoly :: !(SingPoly c)
                   , Root c -> c
rootLbound :: !c
                   , Root c -> c
rootUbound :: !c
                   }
  deriving (Int -> Root c -> ShowS
[Root c] -> ShowS
Root c -> String
(Int -> Root c -> ShowS)
-> (Root c -> String) -> ([Root c] -> ShowS) -> Show (Root c)
forall c. Show c => Int -> Root c -> ShowS
forall c. Show c => [Root c] -> ShowS
forall c. Show c => Root c -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Root c] -> ShowS
$cshowList :: forall c. Show c => [Root c] -> ShowS
show :: Root c -> String
$cshow :: forall c. Show c => Root c -> String
showsPrec :: Int -> Root c -> ShowS
$cshowsPrec :: forall c. Show c => Int -> Root c -> ShowS
Show)

-- | Construct a root from a rational constant
rootFromRational :: Num c => c -> Root c
rootFromRational :: c -> Root c
rootFromRational c
r = Root :: forall c. SingPoly c -> c -> c -> Root c
Root { rootPoly :: SingPoly c
rootPoly = [c] -> SingPoly c
forall c. [c] -> SingPoly c
fromList [ c -> c
forall a. Num a => a -> a
negate c
r, c
1 ]
                          , rootLbound :: c
rootLbound = c
r
                          , rootUbound :: c
rootUbound = c
r
                          }

instance (Ord c, Num c, Pretty c) => Pretty (Root c) where
  pretty :: Root c -> Doc ann
pretty (Root SingPoly c
p c
l c
u) = Doc ann
forall ann. Doc ann
langle Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> SingPoly c -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty SingPoly c
p Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
forall ann. Doc ann
bounds Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
rangle
    where bounds :: Doc ann
bounds = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (c -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty c
l Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> c -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty c
u)

-- | This either returns the root exactly, or it computes the closest double
-- precision approximation of the root.
--
-- Underneath the hood, this uses rational arithmetic to guarantee precision,
-- so this operation is relatively slow.  However, it is guaranteed to provide
-- an exact answer.
--
-- If performance is a concern, there are faster algorithms for computing this.
approximate :: Root Rational -> Rational
approximate :: Root Rational -> Rational
approximate Root Rational
r
    | Rational
l0 Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
u0       = Rational
l0
    | Rational
init_lval Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0 = Rational
l0
    | Rational
init_uval Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0 = Rational
u0
    | Rational
init_lval Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
< Rational
0 Bool -> Bool -> Bool
&& Rational
init_uval Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
> Rational
0 = Double -> Double -> Rational
bisect (Rational -> Double
forall a. Fractional a => Rational -> a
fromRational Rational
l0) (Rational -> Double
forall a. Fractional a => Rational -> a
fromRational Rational
u0)
    | Rational
init_lval Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
> Rational
0 Bool -> Bool -> Bool
&& Rational
init_uval Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
< Rational
0 = Double -> Double -> Rational
bisect (Rational -> Double
forall a. Fractional a => Rational -> a
fromRational Rational
u0) (Rational -> Double
forall a. Fractional a => Rational -> a
fromRational Rational
l0)
    | Bool
otherwise = String -> Rational
forall a. HasCallStack => String -> a
error String
"Closest root given bad root."
  where p_rat :: SingPoly Rational
p_rat = Root Rational -> SingPoly Rational
forall c. Root c -> SingPoly c
rootPoly Root Rational
r
        l0 :: Rational
l0 = Root Rational -> Rational
forall c. Root c -> c
rootLbound Root Rational
r
        u0 :: Rational
u0 = Root Rational -> Rational
forall c. Root c -> c
rootUbound Root Rational
r

        init_lval :: Rational
init_lval = SingPoly Rational -> Rational -> Rational
forall c. Num c => SingPoly c -> c -> c
eval SingPoly Rational
p_rat Rational
l0
        init_uval :: Rational
init_uval = SingPoly Rational -> Rational -> Rational
forall c. Num c => SingPoly c -> c -> c
eval SingPoly Rational
p_rat Rational
u0

        -- bisect takes a value that evaluates to a negative value under the 'p',
        -- and a value that evalautes to a positive value, and runs until it
        -- converges.
        bisect :: Double -> Double -> Rational
        bisect :: Double -> Double -> Rational
bisect Double
l Double
u   -- Stop if mid point is at bound.
                   | Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
l Bool -> Bool -> Bool
|| Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
u = Double -> Rational
forall a. Real a => a -> Rational
toRational (Double -> Rational) -> Double -> Rational
forall a b. (a -> b) -> a -> b
$
                      -- Pick whichever bound is cl oser to root.
                      if Rational
l_val Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
<= Rational
u_val then Double
l else Double
u
                   | Rational
m_val Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0 = Double -> Rational
forall a. Real a => a -> Rational
toRational Double
m -- Stop if mid point is exact root.
                   | Rational
m_val Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
<  Rational
0 = Double -> Double -> Rational
bisect Double
m Double
u -- Use mid point as new lower bound
                   | Bool
otherwise  = Double -> Double -> Rational
bisect Double
l Double
m -- Use mid point as new upper bound.
          where m :: Double
m = (Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
u) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
                m_val :: Rational
m_val = SingPoly Rational -> Rational -> Rational
forall c. Num c => SingPoly c -> c -> c
eval SingPoly Rational
p_rat (Double -> Rational
forall a. Real a => a -> Rational
toRational Double
m)
                l_val :: Rational
l_val = Rational -> Rational
forall a. Num a => a -> a
abs (SingPoly Rational -> Rational -> Rational
forall c. Num c => SingPoly c -> c -> c
eval SingPoly Rational
p_rat (Double -> Rational
forall a. Real a => a -> Rational
toRational Double
l))
                u_val :: Rational
u_val = Rational -> Rational
forall a. Num a => a -> a
abs (SingPoly Rational -> Rational -> Rational
forall c. Num c => SingPoly c -> c -> c
eval SingPoly Rational
p_rat (Double -> Rational
forall a. Real a => a -> Rational
toRational Double
u))


atto_pair :: (a -> b -> r) -> Atto.Parser a -> Atto.Parser b -> Atto.Parser r
atto_pair :: (a -> b -> r) -> Parser a -> Parser b -> Parser r
atto_pair a -> b -> r
f Parser a
x Parser b
y = a -> b -> r
f (a -> b -> r) -> Parser a -> Parser Text (b -> r)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser a
x Parser Text (b -> r) -> Parser b -> Parser r
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (Char -> Parser Char
Atto.char Char
',' Parser Char -> Parser Char -> Parser Char
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Char -> Parser Char
Atto.char Char
' ' Parser Char -> Parser b -> Parser b
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser b
y)

atto_sdecimal :: Integral c => Atto.Parser c
atto_sdecimal :: Parser c
atto_sdecimal = Char -> Parser Char
Atto.char Char
'-' Parser Char -> Parser c -> Parser c
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> (c -> c
forall a. Num a => a -> a
negate (c -> c) -> Parser c -> Parser c
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser c
forall a. Integral a => Parser a
Atto.decimal)
              Parser c -> Parser c -> Parser c
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> Parser c
forall a. Integral a => Parser a
Atto.decimal

atto_rational :: Integral c => Atto.Parser (Ratio c)
atto_rational :: Parser (Ratio c)
atto_rational = c -> c -> Ratio c
forall a. Integral a => a -> a -> Ratio a
(%) (c -> c -> Ratio c) -> Parser Text c -> Parser Text (c -> Ratio c)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Text c
forall a. Integral a => Parser a
atto_sdecimal Parser Text (c -> Ratio c) -> Parser Text c -> Parser (Ratio c)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Parser Text c
denom
  where denom :: Parser Text c
denom = (Char -> Parser Char
Atto.char Char
'/' Parser Char -> Parser Text c -> Parser Text c
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> Parser Text c
forall a. Integral a => Parser a
Atto.decimal) Parser Text c -> Parser Text c -> Parser Text c
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> c -> Parser Text c
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure c
1

parseYicesRoot :: Atto.Parser (Root Rational)
parseYicesRoot :: Parser (Root Rational)
parseYicesRoot = Parser (Root Rational) -> Parser (Root Rational)
forall a. Parser a -> Parser a
atto_angle ((SingPoly Rational -> (Rational, Rational) -> Root Rational)
-> Parser (SingPoly Rational)
-> Parser (Rational, Rational)
-> Parser (Root Rational)
forall a b r. (a -> b -> r) -> Parser a -> Parser b -> Parser r
atto_pair SingPoly Rational -> (Rational, Rational) -> Root Rational
forall c. SingPoly c -> (c, c) -> Root c
mkRoot ((Integer -> Rational) -> SingPoly Integer -> SingPoly Rational
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> Rational
forall a. Num a => Integer -> a
fromInteger (SingPoly Integer -> SingPoly Rational)
-> Parser Text (SingPoly Integer) -> Parser (SingPoly Rational)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Text (SingPoly Integer)
forall c. Integral c => Parser (SingPoly c)
parseYicesPoly) Parser (Rational, Rational)
parseBounds)
             Parser (Root Rational)
-> Parser (Root Rational) -> Parser (Root Rational)
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> (Rational -> Root Rational
forall c. Num c => c -> Root c
rootFromRational (Rational -> Root Rational)
-> Parser Text Rational -> Parser (Root Rational)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Text Rational
forall c. Integral c => Parser (Ratio c)
atto_rational)
  where mkRoot :: SingPoly c -> (c, c) -> Root c
        mkRoot :: SingPoly c -> (c, c) -> Root c
mkRoot = (c -> c -> Root c) -> (c, c) -> Root c
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((c -> c -> Root c) -> (c, c) -> Root c)
-> (SingPoly c -> c -> c -> Root c)
-> SingPoly c
-> (c, c)
-> Root c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SingPoly c -> c -> c -> Root c
forall c. SingPoly c -> c -> c -> Root c
Root
        parseBounds :: Atto.Parser (Rational, Rational)
        parseBounds :: Parser (Rational, Rational)
parseBounds = Parser (Rational, Rational) -> Parser (Rational, Rational)
forall a. Parser a -> Parser a
atto_paren ((Rational -> Rational -> (Rational, Rational))
-> Parser Text Rational
-> Parser Text Rational
-> Parser (Rational, Rational)
forall a b r. (a -> b -> r) -> Parser a -> Parser b -> Parser r
atto_pair (,) Parser Text Rational
forall c. Integral c => Parser (Ratio c)
atto_rational Parser Text Rational
forall c. Integral c => Parser (Ratio c)
atto_rational)

-- | Convert text to a root
fromYicesText :: Text -> Maybe (Root Rational)
fromYicesText :: Text -> Maybe (Root Rational)
fromYicesText Text
t = IResult Text (Root Rational) -> Maybe (Root Rational)
forall a. IResult Text a -> Maybe a
resolve (Parser (Root Rational) -> Text -> IResult Text (Root Rational)
forall a. Parser a -> Text -> Result a
Atto.parse Parser (Root Rational)
parseYicesRoot Text
t)
  where resolve :: IResult Text a -> Maybe a
resolve (Atto.Fail Text
_rem [String]
_ String
_msg) = Maybe a
forall a. Maybe a
Nothing
        resolve (Atto.Partial Text -> IResult Text a
f) =
          IResult Text a -> Maybe a
resolve (Text -> IResult Text a
f Text
Text.empty)
        resolve (Atto.Done Text
i a
r)
          | Text -> Bool
Text.null Text
i = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! a
r
          | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing