From b24f69739e689d8ef15f536fde41eda870d4a5e6 Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Wed, 30 Dec 2020 14:14:13 +0200
Subject: [PATCH] general case for count

---
 vm/src/stdlib/sre.rs        | 21 +++++++++++----------
 vm/src/stdlib/sre/interp.rs | 32 +++++++++++++++++++++++++++-----
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/vm/src/stdlib/sre.rs b/vm/src/stdlib/sre.rs
index 9bec549a1..d8c9ae454 100644
--- a/vm/src/stdlib/sre.rs
+++ b/vm/src/stdlib/sre.rs
@@ -80,13 +80,10 @@ mod _sre {
         indexgroup: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<Pattern> {
-        let code = vm.extract_elements::<u32>(&code)?;
-        dbg!(&code);
-
         Ok(Pattern {
             pattern,
             flags: SreFlag::from_bits_truncate(flags),
-            code,
+            code: vm.extract_elements::<u32>(&code)?,
             groups,
             groupindex,
             indexgroup: vm.extract_elements(&indexgroup)?,
@@ -134,21 +131,21 @@ mod _sre {
     #[pyimpl]
     impl Pattern {
         #[pymethod(name = "match")]
-        fn pymatch(&self, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
+        fn pymatch(zelf: PyRef<Pattern>, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
             interp::pymatch(
                 string_args.string,
                 string_args.pos,
                 string_args.endpos,
-                &self,
+                zelf,
             )
             .map(|x| x.into_ref(vm))
         }
         #[pymethod]
-        fn fullmatch(&self, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
+        fn fullmatch(zelf: PyRef<Pattern>, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
             // TODO: need optimize
             let start = string_args.pos;
             let end = string_args.endpos;
-            let m = self.pymatch(string_args, vm);
+            let m = Self::pymatch(zelf, string_args, vm);
             if let Some(m) = m {
                 if m.start == start && m.end == end {
                     return Some(m);
@@ -157,12 +154,12 @@ mod _sre {
             None
         }
         #[pymethod]
-        fn search(&self, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
+        fn search(zelf: PyRef<Pattern>, string_args: StringArgs, vm: &VirtualMachine) -> Option<PyRef<Match>> {
             // TODO: optimize by op info and skip prefix
             let start = string_args.pos;
             for i in start..string_args.endpos {
                 if let Some(m) =
-                    interp::pymatch(string_args.string.clone(), i, string_args.endpos, &self)
+                    interp::pymatch(string_args.string.clone(), i, string_args.endpos, zelf.clone())
                 {
                     return Some(m.into_ref(vm));
                 }
@@ -189,6 +186,10 @@ mod _sre {
         fn flags(&self) -> u16 {
             self.flags.bits()
         }
+
+        fn subx(&self, sub_args: SubArgs, vm: &VirtualMachine) -> PyResult<PyStrRef> {
+            Err(vm.new_not_implemented_error("".to_owned()))
+        }
     }
 
     #[pyattr]
diff --git a/vm/src/stdlib/sre/interp.rs b/vm/src/stdlib/sre/interp.rs
index 64d70216e..99574443c 100644
--- a/vm/src/stdlib/sre/interp.rs
+++ b/vm/src/stdlib/sre/interp.rs
@@ -3,6 +3,7 @@
 use super::_sre::{Match, Pattern, MAXREPEAT};
 use super::constants::{SreAtCode, SreCatCode, SreFlag, SreOpcode};
 use crate::builtins::PyStrRef;
+use crate::pyobject::PyRef;
 use rustpython_common::borrow::BorrowValue;
 use std::collections::HashMap;
 use std::convert::TryFrom;
@@ -98,7 +99,7 @@ pub(crate) fn pymatch(
     string: PyStrRef,
     start: usize,
     end: usize,
-    pattern: &Pattern,
+    pattern: PyRef<Pattern>,
 ) -> Option<Match> {
     let mut state = State::new(
         string.borrow_value(),
@@ -135,7 +136,7 @@ pub(crate) fn pymatch(
         return None;
     }
 
-    Some(Match::new(&state, pattern.pattern.clone(), string.clone()))
+    Some(Match::new(&state, pattern.clone().into_object(), string.clone()))
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -425,7 +426,7 @@ impl OpcodeDispatcher {
             }),
             SreOpcode::IN_LOC_IGNORE => once(|drive| {
                 let skip = drive.peek_code(1) as usize;
-                if drive.at_end() || !charset_loc_ignore(&drive.pattern()[1..], drive.peek_char()) {
+                if drive.at_end() || !charset_loc_ignore(&drive.pattern()[2..], drive.peek_char()) {
                     drive.ctx_mut().has_matched = Some(false);
                 } else {
                     drive.skip_code(skip + 1);
@@ -561,7 +562,7 @@ fn general_op_literal<F: FnOnce(u32, char) -> bool>(drive: &mut StackDrive, f: F
 
 fn general_op_in<F: FnOnce(char) -> char>(drive: &mut StackDrive, f: F) {
     let skip = drive.peek_code(1) as usize;
-    if drive.at_end() || !charset(&drive.pattern()[1..], f(drive.peek_char())) {
+    if drive.at_end() || !charset(&drive.pattern()[2..], f(drive.peek_char())) {
         drive.ctx_mut().has_matched = Some(false);
     } else {
         drive.skip_code(skip + 1);
@@ -698,7 +699,28 @@ fn charset(set: &[u32], c: char) -> bool {
     false
 }
 
-fn count(stack_drive: &StackDrive, maxcount: usize) -> usize {
+fn count(drive: &mut StackDrive, maxcount: usize) -> usize {
+    let mut count = 0;
+    let maxcount = std::cmp::min(maxcount, drive.remaining_chars());
+
+    let save_ctx = *drive.ctx();
+    drive.skip_code(4);
+    let reset_position = drive.ctx().code_position;
+
+    let mut dispatcher = OpcodeDispatcher::new();
+    while count < maxcount {
+        drive.ctx_mut().code_position = reset_position;
+        dispatcher.dispatch(SreOpcode::try_from(drive.peek_code(0)).unwrap(), drive);
+        if drive.ctx().has_matched == Some(false) {
+            break;
+        }
+        count += 1;
+    }
+    *drive.ctx_mut() = save_ctx;
+    count
+}
+
+fn _count(stack_drive: &StackDrive, maxcount: usize) -> usize {
     let mut drive = WrapDrive::drive(*stack_drive.ctx(), stack_drive);
     let maxcount = std::cmp::min(maxcount, drive.remaining_chars());
     let end = drive.ctx().string_position + maxcount;