From b7a4248837728562d40760ca85ad102ffce13914 Mon Sep 17 00:00:00 2001
From: Nic Ficca <nicficca@gmail.com>
Date: Thu, 7 Feb 2019 19:28:55 -0500
Subject: [PATCH] Add range.count

---
 vm/src/obj/objrange.rs | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs
index 221e98da5..2b6f79617 100644
--- a/vm/src/obj/objrange.rs
+++ b/vm/src/obj/objrange.rs
@@ -66,6 +66,14 @@ impl RangeType {
         }
     }
 
+    #[inline]
+    pub fn count(&self, value: &BigInt) -> usize {
+        match self.index_of(value).is_some() {
+            true => 1,
+            false => 0,
+        }
+    }
+
     #[inline]
     pub fn is_empty(&self) -> bool {
         (self.start <= self.end && self.step.is_negative())
@@ -156,6 +164,7 @@ pub fn init(context: &PyContext) {
         context.new_rustfunc(range_contains),
     );
     context.set_attr(&range_type, "index", context.new_rustfunc(range_index));
+    context.set_attr(&range_type, "count", context.new_rustfunc(range_count));
 }
 
 fn range_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -366,3 +375,20 @@ fn range_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         unreachable!()
     }
 }
+
+fn range_count(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.range_type())), (needle, None)]
+    );
+
+    if let PyObjectPayload::Range { ref range } = zelf.borrow().payload {
+        match needle.borrow().payload {
+            PyObjectPayload::Integer { ref value } => Ok(vm.ctx.new_int(range.count(value))),
+            _ => Ok(vm.ctx.new_int(0)),
+        }
+    } else {
+        unreachable!()
+    }
+}