From 6f499013f6c3984f03e406840cc85fd6bd244888 Mon Sep 17 00:00:00 2001
From: Zhiyan Xiao <zhiyan.xiao@linecorp.com>
Date: Mon, 6 Mar 2023 05:34:18 +0900
Subject: [PATCH] Implement Number Protocol for PyDict

---
 vm/src/builtins/dict.rs | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs
index 5c9c6fcded..dc4a79b33d 100644
--- a/vm/src/builtins/dict.rs
+++ b/vm/src/builtins/dict.rs
@@ -16,11 +16,11 @@ use crate::{
         ArgIterable, FuncArgs, KwArgs, OptionalArg, PyArithmeticValue::*, PyComparisonValue,
     },
     iter::PyExactSizeIterator,
-    protocol::{PyIterIter, PyIterReturn, PyMappingMethods, PySequenceMethods},
+    protocol::{PyIterIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     recursion::ReprGuard,
     types::{
-        AsMapping, AsSequence, Callable, Comparable, Constructor, Hashable, Initializer, IterNext,
-        IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
+        AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor, Hashable, Initializer,
+        IterNext, IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
     },
     vm::VirtualMachine,
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
@@ -209,7 +209,8 @@ impl PyDict {
         Hashable,
         Comparable,
         Iterable,
-        AsSequence
+        AsSequence,
+        AsNumber
     ),
     flags(BASETYPE)
 )]
@@ -477,6 +478,26 @@ impl AsSequence for PyDict {
     }
 }
 
+impl AsNumber for PyDict {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            or: atomic_func!(|num, args, vm| {
+                PyDict::number_downcast(num).or(args.to_pyobject(vm), vm)
+            }),
+            inplace_or: atomic_func!(|num, args, vm| {
+                PyDict::ior(
+                    PyDict::number_downcast(num).to_owned(),
+                    args.to_pyobject(vm),
+                    vm,
+                )
+                .map(|d| d.into())
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl Comparable for PyDict {
     fn cmp(
         zelf: &Py<Self>,