C3 algorithm

chrisrink10 · chrisrink10 · commit 1922c280cfd6 · 2025-06-09T21:19:04.000-04:00
diff --git a/src/singledispatch/core.rs b/src/singledispatch/core.rs
@@ -8,15 +8,16 @@ use pyo3::prelude::*;
 use crate::singledispatch::builtins::Builtins;
 use pyo3::types::{PyDict, PyTuple, PyType};
 use pyo3::{
-    pyclass, pyfunction, pymethods, Bound, IntoPyObjectExt, Py, PyAny, PyObject, PyResult, Python,
+    intern, pyclass, pyfunction, pymethods, Bound, IntoPyObjectExt, Py, PyAny, PyObject, PyResult,
+    Python,
 };
 use std::collections::HashMap;
 use std::sync::Mutex;
 
 fn get_abc_cache_token(py: Python) -> Bound<'_, PyAny> {
-    py.import("abc")
+    py.import(intern!(py, "abc"))
         .unwrap()
-        .getattr("get_cache_token")
+        .getattr(intern!(py, "get_cache_token"))
         .unwrap()
         .call0()
         .unwrap()
@@ -74,8 +75,12 @@ struct SingleDispatchState {
 impl SingleDispatchState {
     fn find_impl(&mut self, py: Python, cls: Bound<'_, PyAny>) -> PyResult<PyObject> {
         let cls_mro = get_obj_mro(&cls.clone()).unwrap();
-        let mro = compose_mro(py, cls.clone(), self.registry.keys())?;
+        let mro = match compose_mro(py, cls.clone(), self.registry.keys()) {
+            Ok(mro) => mro,
+            Err(e) => return Err(e),
+        };
         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));
@@ -95,31 +100,50 @@ impl SingleDispatchState {
                     )));
                 }
                 mro_match = Some(m.clone_ref(py));
+                eprintln!("MRO match: {m}");
                 break;
             }
         }
-        match mro_match {
+        let impl_fn = match mro_match {
             Some(_) => match self.registry.get(&mro_match.unwrap()) {
-                Some(&ref it) => Ok(it.clone_ref(py)),
-                None => Ok(py.None()),
+                Some(&ref it) => Some(it.clone_ref(py)),
+                None => None,
             },
-            None => Ok(py.None()),
+            None => None,
+        };
+        match impl_fn {
+            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!(
+                        "No dispatch function found for {cls}!"
+                    ))),
+                }
+            }
         }
     }
 
     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)),
             None => {
                 let handler_for_cls = match self.registry.get(&type_reference) {
                     Some(handler) => handler.clone_ref(py),
-                    None => self.find_impl(py, free_cls.bind(py).clone())?,
+                    None => match self.find_impl(py, free_cls.bind(py).clone()) {
+                        Ok(handler) => handler,
+                        Err(e) => return Err(e),
+                    },
                 };
                 self.cache
                     .insert(type_reference, handler_for_cls.clone_ref(py));
+                eprintln!("Found new handler {handler_for_cls}");
                 Ok(handler_for_cls)
             }
         }
@@ -161,7 +185,7 @@ impl SingleDispatch {
                     );
                 }
                 if state.cache_token.is_none() {
-                    if let Ok(_) = unbound_func.getattr(py, "__abstractmethods__") {
+                    if let Ok(_) = unbound_func.getattr(py, intern!(py, "__abstractmethods__")) {
                         state.cache_token = Some(get_abc_cache_token(py).unbind());
                     }
                 }
@@ -178,7 +202,7 @@ impl SingleDispatch {
         cls: Bound<'_, PyAny>,
         func: Bound<'_, PyAny>,
     ) -> PyResult<PyObject> {
-        match func.getattr("__annotations__") {
+        match func.getattr(intern!(_py, "__annotations__")) {
             Ok(_annotations) => Err(PyNotImplementedError::new_err("Oops!")),
             Err(_) => Err(PyTypeError::new_err(
                 format!("Invalid first argument to `register()`: {cls}. Use either `@register(some_class)` or plain `@register` on an annotated function."),
@@ -213,22 +237,24 @@ impl SingleDispatch {
         args: &Bound<'_, PyTuple>,
         kwargs: Option<&Bound<'_, PyDict>>,
     ) -> PyResult<Py<PyAny>> {
-        match obj.getattr("__class__") {
+        eprintln!("Calling");
+        match obj.getattr(intern!(py, "__class__")) {
             Ok(cls) => {
                 let mut all_args = Vec::with_capacity(1 + args.len());
                 all_args.insert(0, obj);
                 all_args.extend(args);
 
                 match self.dispatch(py, cls) {
                     Ok(handler) => handler.call(py, PyTuple::new(py, all_args)?, kwargs),
-                    Err(_) => panic!("no handler for singledispatch"),
+                    Err(e) => Err(e),
                 }
             }
             Err(_) => Err(PyTypeError::new_err("expected __class__ attribute for obj")),
         }
     }
 
     fn dispatch(&self, py: Python<'_>, cls: Bound<'_, PyAny>) -> PyResult<PyObject> {
+        eprintln!("Dispatching");
         match self.lock.lock() {
             Ok(mut state) => {
                 match &state.cache_token {
diff --git a/src/singledispatch/mro.rs b/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,18 +19,192 @@ 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 mro: HashSet<_> = cls
+    let subclasses: HashSet<_> = cls
         .call_method0(intern!(cls.py(), "__subclasses__"))?
         .downcast::<PyTuple>()?
         .iter()
         .map(|item| PyTypeReference::new(item.unbind()))
         .collect();
-    Ok(mro)
+    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;
+        }
+    }
+    match candidate {
+        Some(c) => Some(c.clone_ref(py)),
+        None => None,
+    }
+}
+
+struct C3Mro<'a> {
+    seqs: &'a mut Vec<&'a mut Vec<PyTypeReference>>,
 }
 
-fn c3_mro(py: Python, cls: Bound<'_, PyAny>, abcs: Vec<PyTypeReference>) -> PyResult<Vec<PyTypeReference>> {
-    Ok(abcs)
+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>,
+    abcs: Vec<PyTypeReference>,
+) -> PyResult<Vec<PyTypeReference>> {
+    let bases = match get_obj_bases(cls) {
+        Ok(b) => {
+            if b.len() > 0 {
+                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(
@@ -103,5 +278,5 @@ pub(crate) fn compose_mro(
         }
     });
 
-    c3_mro(py, cls, mro)
+    c3_mro(py, &cls, mro)
 }
