From 58be4a8b09772ca790c0b90a859170709c54021c Mon Sep 17 00:00:00 2001 From: lyken Date: Sun, 25 Aug 2024 20:58:56 +0800 Subject: [PATCH] core: add RustSlice::indices --- nac3core/irrt/irrt/slice.hpp | 21 ++++++++- nac3core/src/codegen/irrt/mod.rs | 37 +++++++++++++++ nac3core/src/codegen/object/slice.rs | 67 +++++++++++++++++++++++++++- 3 files changed, 122 insertions(+), 3 deletions(-) diff --git a/nac3core/irrt/irrt/slice.hpp b/nac3core/irrt/irrt/slice.hpp index 0e56362f..3533d054 100644 --- a/nac3core/irrt/irrt/slice.hpp +++ b/nac3core/irrt/irrt/slice.hpp @@ -180,4 +180,23 @@ template struct Slice return this->indices(length); } }; -} // namespace \ No newline at end of file +} // namespace + +extern "C" +{ + void __nac3_slice_indices_i32(bool start_defined, int32_t start, bool stop_defined, int32_t stop, bool step_defined, + int32_t step, int32_t length, int32_t *range_start, int32_t *range_stop, + int32_t *range_step) + { + slice::indices(start_defined, start, stop_defined, stop, step_defined, step, length, range_start, range_stop, + range_step); + } + + void __nac3_slice_indices_i64(bool start_defined, int64_t start, bool stop_defined, int64_t stop, bool step_defined, + int64_t step, int64_t length, int64_t *range_start, int64_t *range_stop, + int64_t *range_step) + { + slice::indices(start_defined, start, stop_defined, stop, step_defined, step, length, range_start, range_stop, + range_step); + } +} \ No newline at end of file diff --git a/nac3core/src/codegen/irrt/mod.rs b/nac3core/src/codegen/irrt/mod.rs index 5c18cdaf..7976ac78 100644 --- a/nac3core/src/codegen/irrt/mod.rs +++ b/nac3core/src/codegen/irrt/mod.rs @@ -635,6 +635,43 @@ pub fn call_nac3_range_len<'ctx, G: CodeGenerator + ?Sized, N: IntKind<'ctx>>( .returning("range_len", Int(int_kind)) } +#[allow(clippy::too_many_arguments)] +pub fn call_nac3_slice_indices<'ctx, G: CodeGenerator + ?Sized, N: IntKind<'ctx>>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + int_kind: N, + start_defined: Instance<'ctx, Int>, + start: Instance<'ctx, Int>, + stop_defined: Instance<'ctx, Int>, + stop: Instance<'ctx, Int>, + step_defined: Instance<'ctx, Int>, + step: Instance<'ctx, Int>, + length: Instance<'ctx, Int>, + range_start: Instance<'ctx, Ptr>>, + range_stop: Instance<'ctx, Ptr>>, + range_step: Instance<'ctx, Ptr>>, +) -> Instance<'ctx, Int> { + let bit_width = int_kind.get_int_type(generator, ctx.ctx).get_bit_width(); + let func_name = match bit_width { + 32 => "__nac3_slice_indices_i32", + 64 => "__nac3_slice_indices_i64", + _ => panic!("{bit_width}-bits ints not supported"), // We could add more variants when necessary. + }; + + CallFunction::begin(generator, ctx, func_name) + .arg(start_defined) + .arg(start) + .arg(stop_defined) + .arg(stop) + .arg(step_defined) + .arg(step) + .arg(length) + .arg(range_start) + .arg(range_stop) + .arg(range_step) + .returning("range_len", Int(int_kind)) +} + pub fn call_nac3_ndarray_util_assert_shape_no_negative<'ctx, G: CodeGenerator + ?Sized>( generator: &mut G, ctx: &mut CodeGenContext<'ctx, '_>, diff --git a/nac3core/src/codegen/object/slice.rs b/nac3core/src/codegen/object/slice.rs index 5b65eea4..7c876ced 100644 --- a/nac3core/src/codegen/object/slice.rs +++ b/nac3core/src/codegen/object/slice.rs @@ -1,4 +1,6 @@ -use crate::codegen::{model::*, CodeGenContext, CodeGenerator}; +use crate::codegen::{irrt::call_nac3_slice_indices, model::*, CodeGenContext, CodeGenerator}; + +use super::range::RustRange; /// Fields of [`Slice`] #[derive(Debug, Clone)] @@ -35,6 +37,9 @@ impl<'ctx, N: IntKind<'ctx>> StructKind<'ctx> for Slice { /// [`Option::None`] if unspecified. #[derive(Debug, Clone)] pub struct RustSlice<'ctx, N: IntKind<'ctx>> { + // It is possible that `start`, `stop`, and `step` are all `None`. + // We need to know the `int_kind` even when that is the case. + pub int_kind: N, pub start: Option>>, pub stop: Option>>, pub step: Option>>, @@ -75,6 +80,64 @@ impl<'ctx, N: IntKind<'ctx>> RustSlice<'ctx, N> { None => dst_slice_ptr.gep(ctx, |f| f.step_defined).store(ctx, false_), } } + + /// Resolve this [`RustSlice`] into a [`RustRange`] like `slice.indices` in Python. + /// + /// NOTE: This function does stack allocation. + pub fn indices( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + length: Instance<'ctx, Int>, + ) -> RustRange<'ctx, N> { + let mut is_defined = |value: Option<_>| -> Instance<'ctx, Int> { + Int(Bool).const_int(generator, ctx.ctx, u64::from(value.is_some())) + }; + + let start_defined = is_defined(self.start); + let stop_defined = is_defined(self.stop); + let step_defined = is_defined(self.step); + + let mut defined_or_zero = |value: Option<_>| -> Instance<'ctx, Int> { + if let Some(value) = value { + value + } else { + // If undefined, return 0 as a placeholder. + Int(self.int_kind).const_0(generator, ctx.ctx) + } + }; + + let start = defined_or_zero(self.start); + let stop = defined_or_zero(self.stop); + let step = defined_or_zero(self.step); + + // Stack allocation here. + let range_start = Int(self.int_kind).alloca(generator, ctx); + let range_stop = Int(self.int_kind).alloca(generator, ctx); + let range_step = Int(self.int_kind).alloca(generator, ctx); + + call_nac3_slice_indices( + generator, + ctx, + self.int_kind, + start_defined, + start, + stop_defined, + stop, + step_defined, + step, + length, + range_start, + range_stop, + range_step, + ); + + let start = range_start.load(generator, ctx); + let stop = range_stop.load(generator, ctx); + let step = range_step.load(generator, ctx); + + RustRange { start, stop, step } + } } pub mod util { @@ -117,6 +180,6 @@ pub mod util { let stop = help(upper)?; let step = help(step)?; - Ok(RustSlice { start, stop, step }) + Ok(RustSlice { int_kind: Int32, start, stop, step }) } }