fix bug; make names a bit more descriptive

plutus-tx/src/PlutusTx/Sqrt.hs

@@ -25,9 +25,14 @@ import qualified Prelude          as Haskell
 
 -- | Integer square-root representation, discarding imaginary integers.
 data Sqrt
+  -- | The number was negative, so we don't even attempt to compute it;
+  -- just note that the result would be imaginary.
   = Imaginary
-  | Exact Integer
-  | Irrational Integer
+  -- | An exact integer result. The 'rsqrt' of 4 is 'Exactly 2'.
+  | Exactly Integer
+  -- | The Integer component of a non-integral result. The 'rsqrt 2' is
+  -- 'Approximately 1'.
+  | Approximately Integer
   deriving stock (Haskell.Show, Haskell.Eq)
 
 {-# INLINABLE rsqrt #-}
@@ -36,17 +41,18 @@ data Sqrt
 rsqrt :: Rational -> Sqrt
 rsqrt r
     | n * d < 0 = Imaginary
-    | n == 0    = Exact 0
-    | n == d    = Exact 1
+    | n == 0    = Exactly 0
+    | n == d    = Exactly 1
     | n < 0     = rsqrt $ negate n % negate d
     | otherwise = go 1 $ 1 + divide n d
   where
     n = numerator r
     d = denominator r
     go :: Integer -> Integer -> Sqrt
     go l u
-        | l * l * d == n = Exact l
-        | u == (l + 1)   = Irrational l
+        | l * l * d == n = Exactly l
+        | u == (l + 1)   = Approximately l
+        | n < d          = Approximately 0
         | otherwise      =
               let
                 m = divide (l + u) 2
@@ -60,7 +66,7 @@ isqrt :: Integer -> Sqrt
 isqrt n = rsqrt (n % 1)
 
 makeLift ''Sqrt
-makeIsDataIndexed ''Sqrt [ ('Imaginary,  0)
-                         , ('Exact,      1)
-                         , ('Irrational, 2)
+makeIsDataIndexed ''Sqrt [ ('Imaginary,     0)
+                         , ('Exactly,       1)
+                         , ('Approximately, 2)
                          ]

plutus-tx/test/Spec.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase       #-}
 {-# LANGUAGE TypeApplications #-}
 module Main(main) where
 
@@ -8,8 +9,9 @@ import           Hedgehog            (MonadGen, Property, annotateShow, assert,
 import qualified Hedgehog.Gen        as Gen
 import qualified Hedgehog.Range      as Range
 import           PlutusTx.Data       (Data (..))
-import           PlutusTx.Ratio      (denominator, numerator, (%))
+import           PlutusTx.Ratio      (Rational, denominator, numerator, (%))
 import           PlutusTx.Sqrt       (Sqrt (..), isqrt, rsqrt)
+import           Prelude             hiding (Rational)
 import           Test.Tasty
 import           Test.Tasty.Hedgehog (testProperty)
 
@@ -25,44 +27,65 @@ tests = testGroup "plutus-tx" [
 sqrtTests :: TestTree
 sqrtTests = testGroup "isqrt/rsqrt tests"
   [ testProperty "isqrt x^2 = x" isqrtRoundTrip
-  , testProperty "rsqrt ((m*x)/x)^2 = m" rsqrtRoundTripExact
+  , testProperty "rsqrt (a/b)^2 = integer part of a/b" rsqrtRoundTrip
+  , testProperty "rsqrt (-x/b) = Imaginary" rsqrtRoundTripImaginary
   ]
 
+rsqrtRoundTripImaginary :: Property
+rsqrtRoundTripImaginary = property $ do
+  let numerators   = Gen.integral (Range.linear (-100000) 0)
+  let denominators = Gen.integral (Range.linear 1 100000)
+
+  -- Note: We're using the fact that (a % -b) is reduced to (-a % b)
+  -- so we only need to test those negative numbers.
+
+  a <- forAll numerators
+  b <- forAll denominators
+
+  let x      = a % b
+      decode = \case
+            Imaginary -> True
+            _         -> False
+
+  assert $ decode (rsqrt x)
+
+rsqrtRoundTrip :: Property
+rsqrtRoundTrip = property $ do
+  let numerators   = Gen.integral (Range.linear 0 100000)
+  let denominators = Gen.integral (Range.linear 1 100000)
+
+  a <- forAll numerators
+  b <- forAll denominators
+
+  let x = a % b
+      f = square
+      g = decode . rsqrt
+      integerPart = a `div` b
+      remainder = rem a b
+      decode = \case
+            Exactly       i -> i == integerPart && remainder == 0
+            Approximately i -> i == integerPart && remainder > 0
+            Imaginary       -> False
+
+  assert $ g (f x)
+
+square :: Rational -> Rational
+square r =
+  let
+    n = numerator r
+    d = denominator r
+    two = 2 :: Integer
+    in (n^two) % (d^two)
+
 isqrtRoundTrip :: Property
 isqrtRoundTrip = property $ do
   let positiveInteger = Gen.integral (Range.linear 0 100000)
   x' <- forAll positiveInteger
   tripping x' sq (decodeExact . isqrt)
     where
       sq x = x ^ (2 :: Integer)
-      decodeExact (Exact x) = Right x
-      decodeExact s         = Left  s
-
-rsqrtRoundTripExact :: Property
-rsqrtRoundTripExact = property $ do
-  let factors      = Gen.integral (Range.linear 0 100000)
-  let denominators = Gen.integral (Range.linear 1 100000)
-
-  n' <- forAll factors
-  d' <- forAll denominators
-
-  let x = (n' * d') % d'
-      f = sq
-      g = decodeExact x . rsqrt
-
-  tripping x f g
-    where
-      sq r =
-        let
-          n = numerator r
-          d = denominator r
-          two = 2 :: Integer
-         in (n^two) % (d^two)
-      decodeExact x s@(Exact e) =
-        if numerator x `div` denominator x == e && rem (numerator x) (denominator x) == 0
-           then Right x
-           else Left s
-      decodeExact _ s         = Left  s
+      decodeExact (Exactly x) = Right x
+      decodeExact s           = Left  s
 
 serdeTests :: TestTree
 serdeTests = testGroup "Data serialisation"

plutus-use-cases/src/Plutus/Contracts/Uniswap.hs

@@ -85,19 +85,19 @@ mkCoin = assetClass
 {-# INLINABLE calculateInitialLiquidity #-}
 calculateInitialLiquidity :: Integer -> Integer -> Integer
 calculateInitialLiquidity outA outB = case isqrt (outA * outB) of
-    Exact l
+    Exactly l
         | l > 0 -> l
-    Irrational l
+    Approximately l
         | l > 0 -> l + 1
     _           -> traceError "insufficient liquidity"
 
 {-# INLINABLE calculateAdditionalLiquidity #-}
 calculateAdditionalLiquidity :: Integer -> Integer -> Integer -> Integer -> Integer -> Integer
 calculateAdditionalLiquidity oldA oldB liquidity delA delB =
   case rsqrt ((liquidity * liquidity * newProd) % oldProd) of
-    Imaginary    -> traceError "insufficient liquidity"
-    Exact x      -> x - liquidity
-    Irrational x -> x - liquidity
+    Imaginary       -> traceError "insufficient liquidity"
+    Exactly x       -> x - liquidity
+    Approximately x -> x - liquidity
   where
     oldProd, newProd :: Integer
     oldProd = oldA * oldB