feat: interoperability of booleans with rationals for if_else (#18)

Author: jcabrero

nada_algebra/__init__.py

@@ -1,5 +1,5 @@
 """This is the __init__.py module"""
 
 from nada_algebra.array import NadaArray
+from nada_algebra.types import RationalConfig, Rational, SecretRational
 from nada_algebra.funcs import *
-from nada_algebra.types import Rational, SecretRational

nada_algebra/types.py

@@ -2,17 +2,21 @@
 
 import numpy as np
 
+import nada_dsl as dsl
+
 from nada_dsl import (
     Input,
     Party,
     UnsignedInteger,
     Integer,
     NadaType,
     SecretInteger,
-    SecretBoolean,
-    PublicBoolean,
+    SecretUnsignedInteger,
     PublicInteger,
+    PublicUnsignedInteger,
 )
+
+
 from typing import Union
 
 
@@ -25,6 +29,85 @@
 
 _NadaRational = Union["Rational", "SecretRational"]
 
+_NadaType = Union[
+    Integer,
+    PublicInteger,
+    PublicUnsignedInteger,
+    SecretInteger,
+    SecretUnsignedInteger,
+    UnsignedInteger,
+]
+
+
+class SecretBoolean(dsl.SecretBoolean):
+
+    def __init__(self, value):
+        super().__init__(value.inner)
+
+    def if_else(
+        self: dsl.SecretBoolean,
+        arg_0: _NadaType | "SecretRational" | "Rational",
+        arg_1: _NadaType | "SecretRational" | "Rational",
+    ) -> Union[SecretInteger, SecretUnsignedInteger]:
+        first_arg = arg_0
+        second_arg = arg_1
+        if isinstance(arg_0, (SecretRational, Rational)) and isinstance(
+            arg_1, (SecretRational, Rational)
+        ):
+            # Both are SecretRational or Rational objects
+            if arg_0.log_scale != arg_1.log_scale:
+                raise ValueError("Cannot output values with different scales.")
+            first_arg = arg_0.value
+            second_arg = arg_1.value
+        elif isinstance(arg_0, (Rational, SecretRational)) or isinstance(
+            arg_1, (Rational, SecretRational)
+        ):
+            # Both are SecretRational or Rational objects
+            raise TypeError(f"Invalid operation: {self}.IfElse({arg_0}, {arg_1})")
+
+        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
+            return SecretRational.from_parts(result, arg_0.log_scale)
+        else:
+            return result
+
+
+class PublicBoolean(dsl.PublicBoolean):
+
+    def __init__(self, value):
+        super().__init__(value.inner)
+
+    def if_else(
+        self: dsl.SecretBoolean,
+        arg_0: _NadaType | "SecretRational" | "Rational",
+        arg_1: _NadaType | "SecretRational" | "Rational",
+    ) -> Union[SecretInteger, SecretUnsignedInteger]:
+        first_arg = arg_0
+        second_arg = arg_1
+        if isinstance(arg_0, (SecretRational, Rational)) and isinstance(
+            arg_1, (SecretRational, Rational)
+        ):
+            # Both are SecretRational or Rational objects
+            if arg_0.log_scale != arg_1.log_scale:
+                raise ValueError("Cannot output values with different scales.")
+            first_arg = arg_0.value
+            second_arg = arg_1.value
+        elif isinstance(arg_0, (Rational, SecretRational)) or isinstance(
+            arg_1, (Rational, SecretRational)
+        ):
+            # Both are SecretRational or Rational objects but of different type
+            raise TypeError(f"Invalid operation: {self}.IfElse({arg_0}, {arg_1})")
+
+        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
+            return Rational.from_parts(result, arg_0.log_scale)
+        else:
+            return result
+
 
 class RationalConfig(object):
 
@@ -828,7 +911,7 @@ def __lt__(self, other: _NadaRational) -> SecretBoolean:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value < other.value
+        return SecretBoolean(self.value < other.value)
 
     def __gt__(self, other: _NadaRational) -> SecretBoolean:
         """Check if this SecretRational is greater than another.
@@ -844,7 +927,7 @@ def __gt__(self, other: _NadaRational) -> SecretBoolean:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value > other.value
+        return SecretBoolean(self.value > other.value)
 
     def __le__(self, other: _NadaRational) -> SecretBoolean:
         """Check if this SecretRational is less than or equal to another.
@@ -860,7 +943,7 @@ def __le__(self, other: _NadaRational) -> SecretBoolean:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value <= other.value
+        return SecretBoolean(self.value <= other.value)
 
     def __ge__(self, other: _NadaRational) -> SecretBoolean:
         """Check if this SecretRational is greater than or equal to another.
@@ -876,7 +959,7 @@ def __ge__(self, other: _NadaRational) -> SecretBoolean:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value >= other.value
+        return SecretBoolean(self.value >= other.value)
 
     def __eq__(self, other: _NadaRational) -> SecretBoolean:
         """Check if this SecretRational is equal to another.
@@ -892,7 +975,7 @@ def __eq__(self, other: _NadaRational) -> SecretBoolean:
         """
         if self.log_scale != other.log_scale:
             raise ValueError("Cannot compare values with different scales.")
-        return self.value == other.value
+        return SecretBoolean(self.value == other.value)
 
     def __ne__(self, other: _NadaRational) -> SecretBoolean:
         """Check if this SecretRational is not equal to another.