X Tutup
Skip to content

Commit 06db3b4

Browse files
youknowoneyouknowone
committed
vm: infer call-convention flags for CPython-style CALL specialization
1 parent 5278f33 commit 06db3b4

File tree

5 files changed

+286
-99
lines changed

5 files changed

+286
-99
lines changed

‎crates/derive-impl/src/pyclass.rs‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
use super::Diagnostic;
22
use crate::util::{
33
ALL_ALLOWED_NAMES, ClassItemMeta, ContentItem, ContentItemInner, ErrorVec, ExceptionItemMeta,
4-
ItemMeta, ItemMetaInner, ItemNursery, SimpleItemMeta, format_doc, pyclass_ident_and_attrs,
5-
pyexception_ident_and_attrs, text_signature,
4+
ItemMeta, ItemMetaInner, ItemNursery, SimpleItemMeta, format_doc, infer_native_call_flags,
5+
pyclass_ident_and_attrs, pyexception_ident_and_attrs, text_signature,
66
};
77
use core::str::FromStr;
88
use proc_macro2::{Delimiter, Group, Span, TokenStream, TokenTree};
@@ -1015,6 +1015,16 @@ where
10151015

10161016
let raw = item_meta.raw()?;
10171017
let sig_doc = text_signature(func.sig(), &py_name);
1018+
let has_receiver = func
1019+
.sig()
1020+
.inputs
1021+
.iter()
1022+
.any(|arg| matches!(arg, syn::FnArg::Receiver(_)));
1023+
let drop_first_typed = match self.inner.attr_name {
1024+
AttrName::Method | AttrName::ClassMethod if !has_receiver => 1,
1025+
_ => 0,
1026+
};
1027+
let call_flags = infer_native_call_flags(func.sig(), drop_first_typed);
10181028

10191029
// Add #[allow(non_snake_case)] for setter methods like set___name__
10201030
let method_name = ident.to_string();
@@ -1031,6 +1041,7 @@ where
10311041
doc,
10321042
raw,
10331043
attr_name: self.inner.attr_name,
1044+
call_flags,
10341045
});
10351046
Ok(())
10361047
}
@@ -1248,6 +1259,7 @@ struct MethodNurseryItem {
12481259
raw: bool,
12491260
doc: Option<String>,
12501261
attr_name: AttrName,
1262+
call_flags: TokenStream,
12511263
}
12521264

12531265
impl MethodNursery {
@@ -1278,7 +1290,7 @@ impl ToTokens for MethodNursery {
12781290
} else {
12791291
quote! { None }
12801292
};
1281-
let flags = match &item.attr_name {
1293+
let binding_flags = match &item.attr_name {
12821294
AttrName::Method => {
12831295
quote! { rustpython_vm::function::PyMethodFlags::METHOD }
12841296
}
@@ -1290,6 +1302,12 @@ impl ToTokens for MethodNursery {
12901302
}
12911303
_ => unreachable!(),
12921304
};
1305+
let call_flags = &item.call_flags;
1306+
let flags = quote! {
1307+
rustpython_vm::function::PyMethodFlags::from_bits_retain(
1308+
(#binding_flags).bits() | (#call_flags).bits()
1309+
)
1310+
};
12931311
// TODO: intern
12941312
// let py_name = if py_name.starts_with("__") && py_name.ends_with("__") {
12951313
// let name_ident = Ident::new(&py_name, ident.span());

‎crates/derive-impl/src/pymodule.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ use crate::error::Diagnostic;
22
use crate::pystructseq::PyStructSequenceMeta;
33
use crate::util::{
44
ALL_ALLOWED_NAMES, AttrItemMeta, AttributeExt, ClassItemMeta, ContentItem, ContentItemInner,
5-
ErrorVec, ItemMeta, ItemNursery, ModuleItemMeta, SimpleItemMeta, format_doc, iter_use_idents,
6-
pyclass_ident_and_attrs, text_signature,
5+
ErrorVec, ItemMeta, ItemNursery, ModuleItemMeta, SimpleItemMeta, format_doc,
6+
infer_native_call_flags, iter_use_idents, pyclass_ident_and_attrs, text_signature,
77
};
88
use core::str::FromStr;
99
use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
@@ -525,6 +525,7 @@ struct FunctionNurseryItem {
525525
cfgs: Vec<Attribute>,
526526
ident: Ident,
527527
doc: String,
528+
call_flags: TokenStream,
528529
}
529530

530531
impl FunctionNursery {
@@ -550,14 +551,14 @@ struct ValidatedFunctionNursery(FunctionNursery);
550551
impl ToTokens for ValidatedFunctionNursery {
551552
fn to_tokens(&self, tokens: &mut TokenStream) {
552553
let mut inner_tokens = TokenStream::new();
553-
let flags = quote! { rustpython_vm::function::PyMethodFlags::empty() };
554554
for item in &self.0.items {
555555
let ident = &item.ident;
556556
let cfgs = &item.cfgs;
557557
let cfgs = quote!(#(#cfgs)*);
558558
let py_names = &item.py_names;
559559
let doc = &item.doc;
560560
let doc = quote!(Some(#doc));
561+
let flags = &item.call_flags;
561562

562563
inner_tokens.extend(quote![
563564
#(
@@ -706,12 +707,14 @@ impl ModuleItem for FunctionItem {
706707
py_names
707708
}
708709
};
710+
let call_flags = infer_native_call_flags(func.sig(), 0);
709711

710712
args.context.function_items.add_item(FunctionNurseryItem {
711713
ident: ident.to_owned(),
712714
py_names,
713715
cfgs: args.cfgs.to_vec(),
714716
doc,
717+
call_flags,
715718
});
716719
Ok(())
717720
}

‎crates/derive-impl/src/util.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -732,6 +732,77 @@ pub(crate) fn text_signature(sig: &Signature, name: &str) -> String {
732732
}
733733
}
734734

735+
pub(crate) fn infer_native_call_flags(sig: &Signature, drop_first_typed: usize) -> TokenStream {
736+
// Best-effort mapping of Rust function signatures to CPython-style
737+
// METH_* calling convention flags used by CALL specialization.
738+
let mut typed_args = Vec::new();
739+
for arg in &sig.inputs {
740+
let syn::FnArg::Typed(typed) = arg else {
741+
continue;
742+
};
743+
let ty_tokens = &typed.ty;
744+
let ty = quote!(#ty_tokens).to_string().replace(' ', "");
745+
// `vm: &VirtualMachine` is not a Python-level argument.
746+
if ty.starts_with('&') && ty.ends_with("VirtualMachine") {
747+
continue;
748+
}
749+
typed_args.push(ty);
750+
}
751+
752+
let mut user_args = typed_args.into_iter();
753+
for _ in 0..drop_first_typed {
754+
if user_args.next().is_none() {
755+
break;
756+
}
757+
}
758+
759+
let mut has_keywords = false;
760+
let mut variable_arity = false;
761+
let mut fixed_positional = 0usize;
762+
763+
for ty in user_args {
764+
let is_named = |name: &str| {
765+
ty == name
766+
|| ty.starts_with(&format!("{name}<"))
767+
|| ty.contains(&format!("::{name}<"))
768+
|| ty.ends_with(&format!("::{name}"))
769+
};
770+
771+
if is_named("FuncArgs") {
772+
has_keywords = true;
773+
variable_arity = true;
774+
continue;
775+
}
776+
if is_named("KwArgs") {
777+
has_keywords = true;
778+
variable_arity = true;
779+
continue;
780+
}
781+
if is_named("PosArgs") || is_named("OptionalArg") || is_named("OptionalOption") {
782+
variable_arity = true;
783+
continue;
784+
}
785+
fixed_positional += 1;
786+
}
787+
788+
if has_keywords {
789+
quote! {
790+
rustpython_vm::function::PyMethodFlags::from_bits_retain(
791+
rustpython_vm::function::PyMethodFlags::FASTCALL.bits()
792+
| rustpython_vm::function::PyMethodFlags::KEYWORDS.bits()
793+
)
794+
}
795+
} else if variable_arity {
796+
quote! { rustpython_vm::function::PyMethodFlags::FASTCALL }
797+
} else {
798+
match fixed_positional {
799+
0 => quote! { rustpython_vm::function::PyMethodFlags::NOARGS },
800+
1 => quote! { rustpython_vm::function::PyMethodFlags::O },
801+
_ => quote! { rustpython_vm::function::PyMethodFlags::FASTCALL },
802+
}
803+
}
804+
}
805+
735806
fn func_sig(sig: &Signature) -> String {
736807
sig.inputs
737808
.iter()

0 commit comments

Comments
 (0)
X Tutup