feat: Add bitwise and, or, xor operations to Nada

NillionNetwork · Sep 5, 2024 · 3fedfd8 · 3fedfd8
1 parent affbf5d
commit 3fedfd8
Show file tree

Hide file tree

Showing 2 changed files with 145 additions and 1 deletion.
diff --git a/nada_dsl/nada_types/types.py b/nada_dsl/nada_types/types.py
@@ -3,7 +3,7 @@
 from . import NadaType
 from dataclasses import dataclass
 from nada_dsl.circuit_io import Literal
-from nada_dsl.operations import Addition, Division, Equals, GreaterOrEqualThan, GreaterThan, IfElse, LeftShift, LessOrEqualThan, LessThan, Modulo, Multiplication, Not, NotEquals, Power, PublicOutputEquality, Random, Reveal, RightShift, Subtraction, TruncPr
+from nada_dsl.operations import Addition, BooleanAnd, BooleanOr, BooleanXor, Division, Equals, GreaterOrEqualThan, GreaterThan, IfElse, LeftShift, LessOrEqualThan, LessThan, Modulo, Multiplication, Not, NotEquals, Power, PublicOutputEquality, Random, Reveal, RightShift, Subtraction, TruncPr
 from nada_dsl.source_ref import SourceRef
 from typing import Union
 
@@ -439,6 +439,48 @@ def __ne__(
         else:
             raise TypeError(f"Invalid operation: {self} != {other}")
 
+    def __and__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["Boolean", "PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            return Boolean(value=bool(self.value & other.value))
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} & {other}")
+
+    def __or__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["Boolean", "PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            return Boolean(value=bool(self.value | other.value))
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} | {other}")
+
+    def __xor__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["Boolean", "PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            return Boolean(value=bool(self.value ^ other.value))
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} ^ {other}")
+
     def __invert__(
         self: "Boolean"
     ) -> "Boolean":
@@ -912,6 +954,51 @@ def __ne__(
         else:
             raise TypeError(f"Invalid operation: {self} != {other}")
 
+    def __and__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} & {other}")
+
+    def __or__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} | {other}")
+
+    def __xor__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> Union["PublicBoolean", "SecretBoolean"]:
+        if isinstance(other, Boolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return PublicBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} ^ {other}")
+
     def __invert__(
         self: "PublicBoolean"
     ) -> "PublicBoolean":
@@ -1467,6 +1554,51 @@ def __ne__(
         else:
             raise TypeError(f"Invalid operation: {self} != {other}")
 
+    def __and__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> "SecretBoolean":
+        if isinstance(other, Boolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanAnd(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} & {other}")
+
+    def __or__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> "SecretBoolean":
+        if isinstance(other, Boolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanOr(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} | {other}")
+
+    def __xor__(
+        self, other: Union["Boolean", "PublicBoolean", "SecretBoolean"]
+    ) -> "SecretBoolean":
+        if isinstance(other, Boolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, PublicBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        elif isinstance(other, SecretBoolean):
+            operation = BooleanXor(left=self, right=other, source_ref=SourceRef.back_frame())
+            return SecretBoolean(inner=operation)
+        else:
+            raise TypeError(f"Invalid operation: {self} ^ {other}")
+
     def __invert__(
         self: "SecretBoolean"
     ) -> "SecretBoolean":

diff --git a/nada_dsl/operations.py b/nada_dsl/operations.py
@@ -109,6 +109,18 @@ class PublicOutputEquality(BinaryOperation):
     pass
 
 
+class BooleanAnd(BinaryOperation):
+    pass
+
+
+class BooleanOr(BinaryOperation):
+    pass
+
+
+class BooleanXor(BinaryOperation):
+    pass
+
+
 class Random:
     source_ref: SourceRef