C3 algorithm

Author: chrisrink10

src/singledispatch/core.rs

@@ -74,6 +74,7 @@ impl SingleDispatchState {
         let cls_mro = get_obj_mro(&cls.clone())?;
         let mro = compose_mro(py, cls.clone(), self.registry.keys())?;
         let mut mro_match: Option<PyTypeReference> = None;
+        eprintln!("Finding impl for {cls}");
         for typ in mro.iter() {
             if self.registry.contains_key(typ) {
                 mro_match = Some(typ.clone_ref(py));
@@ -93,6 +94,7 @@ impl SingleDispatchState {
                     )));
                 }
                 mro_match = Some(m.clone_ref(py));
+                eprintln!("MRO match: {m}");
                 break;
             }
         }
@@ -104,6 +106,7 @@ impl SingleDispatchState {
             Some(f) => Ok(f),
             None => {
                 let obj_type = PyTypeReference::new(Builtins::cached(py).object_type.clone_ref(py));
+                eprintln!("Found impl for {cls}: {obj_type}");
                 match self.registry.get(&obj_type) {
                     Some(it) => Ok(it.clone_ref(py)),
                     None => Err(PyRuntimeError::new_err(format!(
@@ -117,6 +120,7 @@ impl SingleDispatchState {
     fn get_or_find_impl(&mut self, py: Python, cls: Bound<'_, PyAny>) -> PyResult<PyObject> {
         let free_cls = cls.unbind();
         let type_reference = PyTypeReference::new(free_cls.clone_ref(py));
+        eprintln!("Finding impl {type_reference}");
 
         match self.cache.get(&type_reference) {
             Some(handler) => Ok(handler.clone_ref(py)),
@@ -127,6 +131,7 @@ impl SingleDispatchState {
                 };
                 self.cache
                     .insert(type_reference, handler_for_cls.clone_ref(py));
+                eprintln!("Found new handler {handler_for_cls}");
                 Ok(handler_for_cls)
             }
         }
@@ -224,6 +229,7 @@ impl SingleDispatch {
         args: &Bound<'_, PyTuple>,
         kwargs: Option<&Bound<'_, PyDict>>,
     ) -> PyResult<Py<PyAny>> {
+        eprintln!("Calling");
         match obj.getattr(intern!(py, "__class__")) {
             Ok(cls) => {
                 let mut all_args = Vec::with_capacity(1 + args.len());
@@ -240,6 +246,7 @@ impl SingleDispatch {
     }
 
     fn dispatch(&self, py: Python<'_>, cls: Bound<'_, PyAny>) -> PyResult<PyObject> {
+        eprintln!("Dispatching");
         match self.lock.lock() {
             Ok(mut state) => {
                 if let Some(cache_token) = &state.cache_token {

src/singledispatch/mro.rs

@@ -1,6 +1,7 @@
 use crate::singledispatch::builtins::Builtins;
 use crate::singledispatch::typeref::PyTypeReference;
 use crate::singledispatch::typing::TypingModule;
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
 use pyo3::{intern, Bound, PyObject, PyResult, Python};
@@ -18,6 +19,20 @@ pub(crate) fn get_obj_mro(cls: &Bound<'_, PyAny>) -> PyResult<HashSet<PyTypeRefe
     Ok(mro)
 }
 
+fn get_obj_bases(cls: &Bound<'_, PyAny>) -> PyResult<Vec<PyTypeReference>> {
+    match cls.getattr_opt(intern!(cls.py(), "__bases__")) {
+        Ok(opt) => match opt {
+            Some(b) => Ok(b
+                .downcast::<PyTuple>()?
+                .iter()
+                .map(|item| PyTypeReference::new(item.unbind()))
+                .collect()),
+            None => Ok(Vec::new()),
+        },
+        Err(e) => Err(e),
+    }
+}
+
 fn get_obj_subclasses(cls: &Bound<'_, PyAny>) -> PyResult<HashSet<PyTypeReference>> {
     let subclasses: HashSet<_> = cls
         .call_method0(intern!(cls.py(), "__subclasses__"))?
@@ -28,12 +43,165 @@ fn get_obj_subclasses(cls: &Bound<'_, PyAny>) -> PyResult<HashSet<PyTypeReferenc
     Ok(subclasses)
 }
 
+fn find_merge_candidate(py: Python, seqs: &[&mut Vec<PyTypeReference>]) -> Option<PyTypeReference> {
+    let mut candidate: Option<&PyTypeReference> = None;
+    for i1 in 0..seqs.len() {
+        let s1 = &seqs[i1];
+        candidate = Some(&s1[0]);
+        for i2 in 0..seqs.len() {
+            let s2 = &seqs[i2];
+            if s2[1..].contains(candidate.unwrap()) {
+                candidate = None;
+                break;
+            }
+        }
+        if candidate.is_some() {
+            break;
+        }
+    }
+    candidate.map(|c| c.clone_ref(py))
+}
+
+struct C3Mro<'a> {
+    seqs: &'a mut Vec<&'a mut Vec<PyTypeReference>>,
+}
+
+impl C3Mro<'_> {
+    fn for_abcs<'a>(
+        py: Python,
+        abcs: &'a mut Vec<&'a mut Vec<PyTypeReference>>,
+    ) -> PyResult<Vec<PyTypeReference>> {
+        C3Mro { seqs: abcs }.merge(py)
+    }
+
+    fn merge(&mut self, py: Python) -> PyResult<Vec<PyTypeReference>> {
+        let mut result: Vec<PyTypeReference> = Vec::new();
+        loop {
+            let seqs = &mut self.seqs;
+            seqs.retain(|seq| !seq.is_empty());
+            if seqs.is_empty() {
+                return Ok(result);
+            }
+            match find_merge_candidate(py, seqs.as_slice()) {
+                Some(c) => {
+                    for i in 0..seqs.len() {
+                        let seq = &mut self.seqs[i];
+                        if seq[0].eq(&c) {
+                            seq.remove(0);
+                        }
+                    }
+                    result.push(c);
+                }
+                None => return Err(PyRuntimeError::new_err("Inconsistent hierarchy")),
+            }
+        }
+    }
+}
+
+fn c3_boundary(py: Python, bases: &[PyTypeReference]) -> usize {
+    let mut boundary = 0;
+
+    for (i, base) in bases.iter().rev().enumerate() {
+        if base
+            .wrapped()
+            .bind(py)
+            .hasattr(intern!(py, "__abstractmethods__"))
+            .unwrap()
+        {
+            boundary = bases.len() - i;
+            break;
+        }
+    }
+
+    boundary
+}
+
 fn c3_mro(
     py: Python,
-    cls: Bound<'_, PyAny>,
+    cls: &Bound<'_, PyAny>,
     abcs: Vec<PyTypeReference>,
 ) -> PyResult<Vec<PyTypeReference>> {
-    Ok(abcs)
+    let bases = match get_obj_bases(cls) {
+        Ok(b) => {
+            if !b.is_empty() {
+                b
+            } else {
+                return Ok(Vec::new());
+            }
+        }
+        Err(e) => return Err(e),
+    };
+    let boundary = c3_boundary(py, &bases);
+    eprintln!("boundary = {boundary}");
+    let base = &bases[boundary];
+
+    let (explicit_bases, other_bases) = bases.split_at(boundary);
+    let abstract_bases: Vec<_> = abcs
+        .iter()
+        .flat_map(|abc| {
+            if Builtins::cached(py)
+                .issubclass(py, cls, base.wrapped().bind(py))
+                .unwrap()
+                && !bases.iter().any(|b| {
+                    Builtins::cached(py)
+                        .issubclass(py, b.wrapped().bind(py), base.wrapped().bind(py))
+                        .unwrap()
+                })
+            {
+                vec![abc]
+            } else {
+                vec![]
+            }
+        })
+        .collect();
+
+    let new_abcs: Vec<_> = abcs.iter().filter(|c| abstract_bases.contains(c)).collect();
+
+    let mut mros: Vec<&mut Vec<PyTypeReference>> = Vec::new();
+
+    let mut cls_ref = vec![PyTypeReference::new(cls.clone().unbind())];
+    mros.push(&mut cls_ref);
+
+    let mut explicit_bases_mro = Vec::from_iter(explicit_bases.iter().map(|b| {
+        c3_mro(
+            py,
+            b.wrapped().bind(py),
+            new_abcs.iter().map(|abc| abc.clone_ref(py)).collect(),
+        )
+        .unwrap()
+    }));
+    mros.extend(&mut explicit_bases_mro);
+
+    let mut abstract_bases_mro = Vec::from_iter(abstract_bases.iter().map(|b| {
+        c3_mro(
+            py,
+            b.wrapped().bind(py),
+            new_abcs.iter().map(|abc| abc.clone_ref(py)).collect(),
+        )
+        .unwrap()
+    }));
+    mros.extend(&mut abstract_bases_mro);
+
+    let mut other_bases_mro = Vec::from_iter(other_bases.iter().map(|b| {
+        c3_mro(
+            py,
+            b.wrapped().bind(py),
+            new_abcs.iter().map(|abc| abc.clone_ref(py)).collect(),
+        )
+        .unwrap()
+    }));
+    mros.extend(&mut other_bases_mro);
+
+    let mut explicit_bases_cloned = Vec::from_iter(explicit_bases.iter().map(|b| b.clone_ref(py)));
+    mros.push(&mut explicit_bases_cloned);
+
+    let mut abstract_bases_cloned = Vec::from_iter(abstract_bases.iter().map(|b| b.clone_ref(py)));
+    mros.push(&mut abstract_bases_cloned);
+
+    let mut other_bases_cloned = Vec::from_iter(other_bases.iter().map(|b| b.clone_ref(py)));
+    mros.push(&mut other_bases_cloned);
+
+    C3Mro::for_abcs(py, &mut mros)
 }
 
 pub(crate) fn compose_mro(
@@ -107,5 +275,5 @@ pub(crate) fn compose_mro(
         }
     });
 
-    c3_mro(py, cls, mro)
+    c3_mro(py, &cls, mro)
 }