Bugfix/type annotations (#21)

Fix rational if_else logic

Author: mathias-nillion

nada_algebra/__init__.py

@@ -5,6 +5,8 @@
 from nada_algebra.types import (
     Rational,
     SecretRational,
+    PublicBoolean,
+    SecretBoolean,
     public_rational,
     rational,
     secret_rational,

nada_algebra/array.py

@@ -614,3 +614,13 @@ def dtype(self) -> Type:
         if self.empty:
             return NoneType
         return type(self.inner.item(0))
+
+    @property
+    def is_rational(self) -> bool:
+        """
+        Returns whether or not the Array's type is a rational.
+
+        Returns:
+            bool: Boolean output.
+        """
+        return self.dtype in (Rational, SecretRational)

nada_algebra/types.py

@@ -48,7 +48,7 @@ def if_else(
         self,
         arg_0: Union[_NadaType, "SecretRational", "Rational"],
         arg_1: Union[_NadaType, "SecretRational", "Rational"],
-    ) -> Union[SecretInteger, SecretUnsignedInteger]:
+    ) -> Union[SecretInteger, SecretUnsignedInteger, "SecretRational"]:
         """
         If-else logic. If the boolean is True, arg_0 is returned. If not, arg_1 is returned.
 
@@ -61,7 +61,7 @@ def if_else(
             TypeError: Raised when invalid operation is called.
 
         Returns:
-            Union[SecretInteger, SecretUnsignedInteger]: Return value.
+            Union[SecretInteger, SecretUnsignedInteger, "SecretRational"]: Return value.
         """
         first_arg = arg_0
         second_arg = arg_1
@@ -84,8 +84,7 @@ def if_else(
         if isinstance(arg_0, (SecretRational, Rational)):
             # If we have a SecretBoolean, the return type will be SecretInteger, thus promoted to SecretRational
             return SecretRational(result, arg_0.log_scale, is_scaled=True)
-        else:
-            return result
+        return result
 
 
 class PublicBoolean(dsl.PublicBoolean):
@@ -104,7 +103,14 @@ def if_else(
         self,
         arg_0: Union[_NadaType, "SecretRational", "Rational"],
         arg_1: Union[_NadaType, "SecretRational", "Rational"],
-    ) -> Union[SecretInteger, SecretUnsignedInteger]:
+    ) -> Union[
+        PublicInteger,
+        PublicUnsignedInteger,
+        SecretInteger,
+        SecretUnsignedInteger,
+        "Rational",
+        "SecretRational",
+    ]:
         """
         If-else logic. If the boolean is True, arg_0 is returned. If not, arg_1 is returned.
 
@@ -117,7 +123,8 @@ def if_else(
             TypeError: Raised when invalid operation is called.
 
         Returns:
-            Union[SecretInteger, SecretUnsignedInteger]: Return value.
+            Union[PublicInteger, PublicUnsignedInteger, SecretInteger,
+                SecretUnsignedInteger, "Rational", "SecretRational"]: Return value.
         """
         first_arg = arg_0
         second_arg = arg_1
@@ -137,11 +144,11 @@ def if_else(
 
         result = super().if_else(first_arg, second_arg)
 
-        if isinstance(arg_0, (SecretRational, Rational)):
-            # If we have a SecretBoolean, the return type will be SecretInteger, thus promoted to SecretRational
+        if isinstance(arg_0, SecretRational) or isinstance(arg_1, SecretRational):
+            return SecretRational(result, arg_0.log_scale, is_scaled=True)
+        elif isinstance(arg_0, Rational) and isinstance(arg_1, Rational):
             return Rational(result, arg_0.log_scale, is_scaled=True)
-        else:
-            return result
+        return result
 
 
 class Rational:
@@ -545,7 +552,9 @@ def __lt__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value < other.value
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value < other.value)
+        return PublicBoolean(self.value < other.value)
 
     def __gt__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
@@ -562,7 +571,9 @@ def __gt__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value > other.value
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value > other.value)
+        return PublicBoolean(self.value > other.value)
 
     def __le__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
@@ -579,7 +590,9 @@ def __le__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value <= other.value
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value <= other.value)
+        return PublicBoolean(self.value <= other.value)
 
     def __ge__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
@@ -596,7 +609,9 @@ def __ge__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value >= other.value
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value >= other.value)
+        return PublicBoolean(self.value >= other.value)
 
     def __eq__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
@@ -613,7 +628,9 @@ def __eq__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value == other.value
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value == other.value)
+        return PublicBoolean(self.value == other.value)
 
     def __ne__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
@@ -630,7 +647,9 @@ def __ne__(self, other: _NadaRational) -> Union[PublicBoolean, SecretBoolean]:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return SecretBoolean(self.value != other.value)
+        if isinstance(other, SecretRational):
+            return SecretBoolean(self.value != other.value)
+        return PublicBoolean(self.value != other.value)
 
     def rescale_up(self, log_scale: int = None) -> "Rational":
         """
@@ -1287,6 +1306,9 @@ def rational(
     Returns:
         Rational: Instantiated Rational object.
     """
+    if value == 0:  # no use in rescaling 0
+        return Rational(Integer(0), is_scaled=True)
+
     if log_scale is None:
         log_scale = get_log_scale()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nada-algebra"
-version = "0.3.0"
+version = "0.3.1"
 description = "Nada-Algebra is a Python library designed for algebraic operations on NumPy-like array objects on top of Nada DSL and Nillion Network."
 authors = ["José Cabrero-Holgueras <jose.cabrero@nillion.com>"]
 readme = "README.md"

tests/nada-tests/src/rational_if_else.py

@@ -1,3 +1,4 @@
+import pytest
 from nada_dsl import *
 import nada_algebra as na
 
@@ -8,16 +9,77 @@ def nada_main():
     a = na.secret_rational("A", parties[0])
     b = na.secret_rational("B", parties[1])
     c = na.secret_rational("C", parties[2])
+    d = SecretInteger(Input("D", parties[0]))
+    e = SecretUnsignedInteger(Input("E", parties[0]))
 
-    out_0 = (a > b).if_else(na.rational(0), na.rational(1))
-    out_1 = (a >= b).if_else(na.rational(1), na.rational(0))
-    out_2 = (a < b).if_else(na.rational(2), na.rational(1))
-    out_3 = (a <= b).if_else(na.rational(3), na.rational(0))
+    out_0 = (a > b).if_else(a, b)
+    assert isinstance(out_0, na.SecretRational), type(out_0).__name__
+    out_1 = (a >= b).if_else(a, b)
+    assert isinstance(out_1, na.SecretRational), type(out_1).__name__
+    out_2 = (a < b).if_else(a, b)
+    assert isinstance(out_2, na.SecretRational), type(out_2).__name__
+    out_3 = (a <= b).if_else(a, b)
+    assert isinstance(out_3, na.SecretRational), type(out_3).__name__
 
-    out_4 = (a > b).if_else(c, na.rational(2))
-    out_5 = (a >= b).if_else(na.rational(2), c)
-    out_6 = (a < b).if_else(c, na.rational(2))
-    out_7 = (a <= b).if_else(c, na.rational(2))
+    out_4 = (a > na.rational(1)).if_else(d, Integer(1))
+    assert isinstance(out_4, SecretInteger), type(out_4).__name__
+    out_5 = (a >= na.rational(1)).if_else(d, Integer(1))
+    assert isinstance(out_5, SecretInteger), type(out_5).__name__
+    out_6 = (a < na.rational(1)).if_else(d, Integer(1))
+    assert isinstance(out_6, SecretInteger), type(out_6).__name__
+    out_7 = (a <= na.rational(1)).if_else(d, Integer(1))
+    assert isinstance(out_7, SecretInteger), type(out_7).__name__
+
+    out_8 = (na.rational(0) > na.rational(1)).if_else(na.rational(1), na.rational(2))
+    assert isinstance(out_8, na.Rational), type(out_8).__name__
+    out_9 = (na.rational(0) >= na.rational(1)).if_else(na.rational(2), na.rational(1))
+    assert isinstance(out_9, na.Rational), type(out_9).__name__
+    out_10 = (na.rational(0) < na.rational(1)).if_else(Integer(1), Integer(2))
+    assert isinstance(out_10, PublicInteger), type(out_10).__name__
+    out_11 = (na.rational(0) <= na.rational(1)).if_else(Integer(1), d)
+    assert isinstance(out_11, SecretInteger), type(out_11).__name__
+    out_12 = (na.rational(0) <= na.rational(1)).if_else(UnsignedInteger(1), e)
+    assert isinstance(out_12, SecretUnsignedInteger), type(out_12).__name__
+    out_13 = (na.rational(0) <= na.rational(1)).if_else(
+        UnsignedInteger(1), UnsignedInteger(0)
+    )
+    assert isinstance(out_13, PublicUnsignedInteger), type(out_13).__name__
+
+    # Incompatible input types
+    with pytest.raises(Exception):
+        (a > Integer(1)).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (Integer(1) > a).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (a > d).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (d > a).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (na.rational(1) > Integer(1)).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (Integer(1) > na.rational(1)).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (na.rational(1) > d).if_else(na.rational(0), na.rational(1))
+    with pytest.raises(Exception):
+        (d > na.rational(1)).if_else(na.rational(0), na.rational(1))
+
+    # Incompatible return types
+    with pytest.raises(Exception):
+        (a > b).if_else(c, Integer(1))
+    with pytest.raises(Exception):
+        (a > b).if_else(Integer(1), c)
+    with pytest.raises(Exception):
+        (a > b).if_else(c, d)
+    with pytest.raises(Exception):
+        (a > b).if_else(d, c)
+    with pytest.raises(Exception):
+        (na.rational(0) > na.rational(1)).if_else(c, Integer(1))
+    with pytest.raises(Exception):
+        (na.rational(0) > na.rational(1)).if_else(Integer(1), c)
+    with pytest.raises(Exception):
+        (na.rational(0) > na.rational(1)).if_else(c, d)
+    with pytest.raises(Exception):
+        (na.rational(0) > na.rational(1)).if_else(d, c)
 
     return (
         na.output(out_0, parties[2], "out_0")
@@ -28,4 +90,10 @@ def nada_main():
         + na.output(out_5, parties[2], "out_5")
         + na.output(out_6, parties[2], "out_6")
         + na.output(out_7, parties[2], "out_7")
+        + na.output(out_8, parties[2], "out_8")
+        + na.output(out_9, parties[2], "out_9")
+        + na.output(out_10, parties[2], "out_10")
+        + na.output(out_11, parties[2], "out_11")
+        + na.output(out_12, parties[2], "out_12")
+        + na.output(out_13, parties[2], "out_13")
     )

tests/nada-tests/tests/rational_if_else.yaml

@@ -8,21 +8,35 @@ inputs:
       SecretInteger: "294912" # 3.2
     C:
       SecretInteger: "65536" # 1
+    D:
+      SecretInteger: "1"
+    E:
+      SecretUnsignedInteger: "1"
   public_variables: {}
 expected_outputs:
   out_0_0:
-    SecretInteger: "65536"
+    SecretInteger: "294912"
   out_1_0:
-    SecretInteger: "0"
+    SecretInteger: "294912"
   out_2_0:
-    SecretInteger: "131072"
+    SecretInteger: "78643"
   out_3_0:
-    SecretInteger: "196608"
+    SecretInteger: "78643"
   out_4_0:
-    SecretInteger: "131072"
+    SecretInteger: "1"
   out_5_0:
-    SecretInteger: "65536"
+    SecretInteger: "1"
   out_6_0:
-    SecretInteger: "65536"
+    SecretInteger: "1"
   out_7_0:
-    SecretInteger: "65536"
+    SecretInteger: "1"
+  out_8_0:
+    Integer: "131072"
+  out_9_0:
+    Integer: "65536"
+  out_10_0:
+    Integer: "1"
+  out_11_0:
+    SecretInteger: "1"
+  out_12_0:
+    SecretUnsignedInteger: "1"