1
0
forked from M-Labs/nac3

core/classes: Take reference of indexes

This commit is contained in:
David Mak 2024-04-16 17:20:24 +08:00
parent 025b3cd02f
commit 5dfcc63978
6 changed files with 67 additions and 64 deletions

View File

@ -58,7 +58,7 @@ pub trait ArrayLikeIndexer<'ctx, Index = IntValue<'ctx>>: ArrayLikeValue<'ctx> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> PointerValue<'ctx>;
@ -67,7 +67,7 @@ pub trait ArrayLikeIndexer<'ctx, Index = IntValue<'ctx>>: ArrayLikeValue<'ctx> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> PointerValue<'ctx>;
}
@ -81,7 +81,7 @@ pub trait UntypedArrayLikeAccessor<'ctx, Index = IntValue<'ctx>>: ArrayLikeIndex
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> BasicValueEnum<'ctx> {
let ptr = self.ptr_offset_unchecked(ctx, generator, idx, name);
@ -93,7 +93,7 @@ pub trait UntypedArrayLikeAccessor<'ctx, Index = IntValue<'ctx>>: ArrayLikeIndex
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> BasicValueEnum<'ctx> {
let ptr = self.ptr_offset(ctx, generator, idx, name);
@ -110,7 +110,7 @@ pub trait UntypedArrayLikeMutator<'ctx, Index = IntValue<'ctx>>: ArrayLikeIndexe
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
value: BasicValueEnum<'ctx>,
) {
let ptr = self.ptr_offset_unchecked(ctx, generator, idx, None);
@ -122,7 +122,7 @@ pub trait UntypedArrayLikeMutator<'ctx, Index = IntValue<'ctx>>: ArrayLikeIndexe
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
value: BasicValueEnum<'ctx>,
) {
let ptr = self.ptr_offset(ctx, generator, idx, None);
@ -142,7 +142,7 @@ pub trait TypedArrayLikeAccessor<'ctx, T, Index = IntValue<'ctx>>: UntypedArrayL
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> T {
let value = self.get_unchecked(ctx, generator, idx, name);
@ -154,7 +154,7 @@ pub trait TypedArrayLikeAccessor<'ctx, T, Index = IntValue<'ctx>>: UntypedArrayL
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> T {
let value = self.get(ctx, generator, idx, name);
@ -174,7 +174,7 @@ pub trait TypedArrayLikeMutator<'ctx, T, Index = IntValue<'ctx>>: UntypedArrayLi
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
value: T,
) {
let value = self.upcast_from_type(ctx, value);
@ -186,7 +186,7 @@ pub trait TypedArrayLikeMutator<'ctx, T, Index = IntValue<'ctx>>: UntypedArrayLi
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
value: T,
) {
let value = self.upcast_from_type(ctx, value);
@ -255,7 +255,7 @@ impl<'ctx, T, Index, Adapted> ArrayLikeIndexer<'ctx, Index> for TypedArrayLikeAd
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> PointerValue<'ctx> {
self.adapted.ptr_offset_unchecked(ctx, generator, idx, name)
@ -265,7 +265,7 @@ impl<'ctx, T, Index, Adapted> ArrayLikeIndexer<'ctx, Index> for TypedArrayLikeAd
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: Index,
idx: &Index,
name: Option<&str>,
) -> PointerValue<'ctx> {
self.adapted.ptr_offset(ctx, generator, idx, name)
@ -345,7 +345,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ArraySliceValue<'ctx> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let var_name = name
@ -354,7 +354,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ArraySliceValue<'ctx> {
ctx.builder.build_in_bounds_gep(
self.base_ptr(ctx, generator),
&[idx],
&[*idx],
var_name.as_str(),
).unwrap()
}
@ -363,13 +363,13 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ArraySliceValue<'ctx> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
debug_assert_eq!(idx.get_type(), generator.get_size_type(ctx.ctx));
let size = self.size(ctx, generator);
let in_range = ctx.builder.build_int_compare(IntPredicate::ULT, idx, size, "").unwrap();
let in_range = ctx.builder.build_int_compare(IntPredicate::ULT, *idx, size, "").unwrap();
ctx.make_assert(
generator,
in_range,
@ -573,7 +573,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ListDataProxy<'ctx, '_> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let var_name = name
@ -582,7 +582,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ListDataProxy<'ctx, '_> {
ctx.builder.build_in_bounds_gep(
self.base_ptr(ctx, generator),
&[idx],
&[*idx],
var_name.as_str(),
).unwrap()
}
@ -591,13 +591,13 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for ListDataProxy<'ctx, '_> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
debug_assert_eq!(idx.get_type(), generator.get_size_type(ctx.ctx));
let size = self.size(ctx, generator);
let in_range = ctx.builder.build_int_compare(IntPredicate::ULT, idx, size, "").unwrap();
let in_range = ctx.builder.build_int_compare(IntPredicate::ULT, *idx, size, "").unwrap();
ctx.make_assert(
generator,
in_range,
@ -1015,7 +1015,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx, IntValue<'ctx>> for NDArrayDimsProxy<'ctx, '_>
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let var_name = name
@ -1024,7 +1024,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx, IntValue<'ctx>> for NDArrayDimsProxy<'ctx, '_>
ctx.builder.build_in_bounds_gep(
self.base_ptr(ctx, generator),
&[idx],
&[*idx],
var_name.as_str(),
).unwrap()
}
@ -1033,13 +1033,13 @@ impl<'ctx> ArrayLikeIndexer<'ctx, IntValue<'ctx>> for NDArrayDimsProxy<'ctx, '_>
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let size = self.size(ctx, generator);
let in_range = ctx.builder.build_int_compare(
IntPredicate::ULT,
idx,
*idx,
size,
""
).unwrap();
@ -1048,7 +1048,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx, IntValue<'ctx>> for NDArrayDimsProxy<'ctx, '_>
in_range,
"0:IndexError",
"index {0} is out of bounds for axis 0 with size {1}",
[Some(idx), Some(self.0.load_ndims(ctx)), None],
[Some(*idx), Some(self.0.load_ndims(ctx)), None],
ctx.current_loc,
);
@ -1120,12 +1120,12 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for NDArrayDataProxy<'ctx, '_> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
ctx.builder.build_in_bounds_gep(
self.base_ptr(ctx, generator),
&[idx],
&[*idx],
name.unwrap_or_default(),
).unwrap()
}
@ -1134,13 +1134,13 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for NDArrayDataProxy<'ctx, '_> {
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
idx: IntValue<'ctx>,
idx: &IntValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let data_sz = self.size(ctx, generator);
let in_range = ctx.builder.build_int_compare(
IntPredicate::ULT,
idx,
*idx,
data_sz,
""
).unwrap();
@ -1149,7 +1149,7 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for NDArrayDataProxy<'ctx, '_> {
in_range,
"0:IndexError",
"index {0} is out of bounds with size {1}",
[Some(idx), Some(self.0.load_ndims(ctx)), None],
[Some(*idx), Some(self.0.load_ndims(ctx)), None],
ctx.current_loc,
);
@ -1167,12 +1167,15 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
indices: Index,
indices: &Index,
name: Option<&str>,
) -> PointerValue<'ctx> {
let llvm_usize = generator.get_size_type(ctx.ctx);
let indices_elem_ty = indices.ptr_offset(ctx, generator, llvm_usize.const_zero(), None).get_type().get_element_type();
let indices_elem_ty = indices
.ptr_offset(ctx, generator, &llvm_usize.const_zero(), None)
.get_type()
.get_element_type();
let Ok(indices_elem_ty) = IntType::try_from(indices_elem_ty) else {
panic!("Expected list[int32] but got {indices_elem_ty}")
};
@ -1182,7 +1185,7 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
generator,
ctx,
*self.0,
&indices,
indices,
);
unsafe {
@ -1198,7 +1201,7 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
indices: Index,
indices: &Index,
name: Option<&str>,
) -> PointerValue<'ctx> {
let llvm_usize = generator.get_size_type(ctx.ctx);
@ -1230,8 +1233,8 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
|generator, ctx, i| {
let (dim_idx, dim_sz) = unsafe {
(
indices.get_unchecked(ctx, generator, i, None).into_int_value(),
self.0.dim_sizes().get_typed_unchecked(ctx, generator, i, None),
indices.get_unchecked(ctx, generator, &i, None).into_int_value(),
self.0.dim_sizes().get_typed_unchecked(ctx, generator, &i, None),
)
};
let dim_idx = ctx.builder

View File

@ -1719,7 +1719,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
v.dim_sizes().get_typed_unchecked(
ctx,
generator,
llvm_usize.const_zero(),
&llvm_usize.const_zero(),
None,
)
};
@ -1746,7 +1746,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
.get(
ctx,
generator,
ArraySliceValue::from_ptr_val(index_addr, llvm_usize.const_int(1, false), None),
&ArraySliceValue::from_ptr_val(index_addr, llvm_usize.const_int(1, false), None),
None,
)
.into()))
@ -1781,7 +1781,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
v.dim_sizes().ptr_offset_unchecked(
ctx,
generator,
llvm_usize.const_int(1, false),
&llvm_usize.const_int(1, false),
None,
)
};
@ -1806,7 +1806,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
let v_data_src_ptr = v.data().ptr_offset(
ctx,
generator,
ArraySliceValue::from_ptr_val(index_addr, llvm_usize.const_int(1, false), None),
&ArraySliceValue::from_ptr_val(index_addr, llvm_usize.const_int(1, false), None),
None
);
call_memcpy_generic(
@ -1906,7 +1906,7 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
let arr_ptr = arr_str_ptr.data();
for (i, v) in elements.iter().enumerate() {
let elem_ptr = arr_ptr
.ptr_offset(ctx, generator, usize.const_int(i as u64, false), Some("elem_ptr"));
.ptr_offset(ctx, generator, &usize.const_int(i as u64, false), Some("elem_ptr"));
ctx.builder.build_store(elem_ptr, *v).unwrap();
}
arr_str_ptr.as_ptr_value().into()
@ -2324,7 +2324,7 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
[Some(raw_index), Some(len), None],
expr.location,
);
v.data().get(ctx, generator, index, None).into()
v.data().get(ctx, generator, &index, None).into()
}
}
TypeEnum::TObj { obj_id, params, .. } if *obj_id == PRIMITIVE_DEF_IDS.ndarray => {

View File

@ -833,8 +833,8 @@ pub fn call_ndarray_calc_broadcast<'ctx, G: CodeGenerator + ?Sized>(
let idx = ctx.builder.build_int_sub(min_ndims, idx, "").unwrap();
let (lhs_dim_sz, rhs_dim_sz) = unsafe {
(
lhs.dim_sizes().get_typed_unchecked(ctx, generator, idx, None),
rhs.dim_sizes().get_typed_unchecked(ctx, generator, idx, None),
lhs.dim_sizes().get_typed_unchecked(ctx, generator, &idx, None),
rhs.dim_sizes().get_typed_unchecked(ctx, generator, &idx, None),
)
};
@ -955,7 +955,7 @@ pub fn call_ndarray_calc_broadcast_index<'ctx, G: CodeGenerator + ?Sized, Broadc
broadcast_idx.ptr_offset_unchecked(
ctx,
generator,
llvm_usize.const_zero(),
&llvm_usize.const_zero(),
None
)
};

View File

@ -119,7 +119,7 @@ fn create_ndarray_dyn_shape<'ctx, 'a, G, V, LenFn, DataFn>(
.unwrap();
let ndarray_pdim = unsafe {
ndarray.dim_sizes().ptr_offset_unchecked(ctx, generator, i, None)
ndarray.dim_sizes().ptr_offset_unchecked(ctx, generator, &i, None)
};
ctx.builder.build_store(ndarray_pdim, shape_dim).unwrap();
@ -190,7 +190,7 @@ fn create_ndarray_const_shape<'ctx, G: CodeGenerator + ?Sized>(
let ndarray_dim = unsafe {
ndarray
.dim_sizes()
.ptr_offset_unchecked(ctx, generator, llvm_usize.const_int(i as u64, true), None)
.ptr_offset_unchecked(ctx, generator, &llvm_usize.const_int(i as u64, true), None)
};
ctx.builder.build_store(ndarray_dim, *shape_dim).unwrap();
@ -267,7 +267,7 @@ fn call_ndarray_empty_impl<'ctx, G: CodeGenerator + ?Sized>(
Ok(shape.load_size(ctx, None))
},
|generator, ctx, shape, idx| {
Ok(shape.data().get(ctx, generator, idx, None).into_int_value())
Ok(shape.data().get(ctx, generator, &idx, None).into_int_value())
},
)
}
@ -299,7 +299,7 @@ fn ndarray_fill_flattened<'ctx, 'a, G, ValueFn>(
(ndarray_num_elems, false),
|generator, ctx, i| {
let elem = unsafe {
ndarray.data().ptr_offset_unchecked(ctx, generator, i, None)
ndarray.data().ptr_offset_unchecked(ctx, generator, &i, None)
};
let value = value_fn(generator, ctx, i)?;
@ -321,7 +321,7 @@ fn ndarray_fill_indexed<'ctx, G, ValueFn>(
) -> Result<(), String>
where
G: CodeGenerator + ?Sized,
ValueFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, TypedArrayLikeAdapter<'ctx, IntValue<'ctx>>) -> Result<BasicValueEnum<'ctx>, String>,
ValueFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, &TypedArrayLikeAdapter<'ctx, IntValue<'ctx>>) -> Result<BasicValueEnum<'ctx>, String>,
{
ndarray_fill_flattened(
generator,
@ -335,7 +335,7 @@ fn ndarray_fill_indexed<'ctx, G, ValueFn>(
ndarray,
);
value_fn(generator, ctx, indices)
value_fn(generator, ctx, &indices)
}
)
}
@ -357,7 +357,7 @@ fn ndarray_fill_mapping<'ctx, G, MapFn>(
dest,
|generator, ctx, i| {
let elem = unsafe {
src.data().get_unchecked(ctx, generator, i, None)
src.data().get_unchecked(ctx, generator, &i, None)
};
map_fn(generator, ctx, elem)
@ -430,10 +430,10 @@ fn ndarray_broadcast_fill<'ctx, G, ValueFn>(
lhs_val
} else {
let lhs = NDArrayValue::from_ptr_val(lhs_val.into_pointer_value(), llvm_usize, None);
let lhs_idx = call_ndarray_calc_broadcast_index(generator, ctx, lhs, &idx);
let lhs_idx = call_ndarray_calc_broadcast_index(generator, ctx, lhs, idx);
unsafe {
lhs.data().get_unchecked(ctx, generator, lhs_idx, None)
lhs.data().get_unchecked(ctx, generator, &lhs_idx, None)
}
};
@ -441,10 +441,10 @@ fn ndarray_broadcast_fill<'ctx, G, ValueFn>(
rhs_val
} else {
let rhs = NDArrayValue::from_ptr_val(rhs_val.into_pointer_value(), llvm_usize, None);
let rhs_idx = call_ndarray_calc_broadcast_index(generator, ctx, rhs, &idx);
let rhs_idx = call_ndarray_calc_broadcast_index(generator, ctx, rhs, idx);
unsafe {
rhs.data().get_unchecked(ctx, generator, rhs_idx, None)
rhs.data().get_unchecked(ctx, generator, &rhs_idx, None)
}
};
@ -604,8 +604,8 @@ fn call_ndarray_eye_impl<'ctx, G: CodeGenerator + ?Sized>(
|generator, ctx, indices| {
let (row, col) = unsafe {
(
indices.get_typed_unchecked(ctx, generator, llvm_usize.const_zero(), None),
indices.get_typed_unchecked(ctx, generator, llvm_usize.const_int(1, false), None),
indices.get_typed_unchecked(ctx, generator, &llvm_usize.const_zero(), None),
indices.get_typed_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None),
)
};
@ -652,7 +652,7 @@ fn ndarray_copy_impl<'ctx, G: CodeGenerator + ?Sized>(
Ok(shape.load_ndims(ctx))
},
|generator, ctx, shape, idx| {
unsafe { Ok(shape.dim_sizes().get_typed_unchecked(ctx, generator, idx, None)) }
unsafe { Ok(shape.dim_sizes().get_typed_unchecked(ctx, generator, &idx, None)) }
},
)?;
@ -704,7 +704,7 @@ pub fn ndarray_elementwise_unaryop_impl<'ctx, G, MapFn>(
},
|generator, ctx, v, idx| {
unsafe {
Ok(v.dim_sizes().get_typed_unchecked(ctx, generator, idx, None))
Ok(v.dim_sizes().get_typed_unchecked(ctx, generator, &idx, None))
}
},
).unwrap()
@ -782,7 +782,7 @@ pub fn ndarray_elementwise_binop_impl<'ctx, G, ValueFn>(
},
|generator, ctx, v, idx| {
unsafe {
Ok(v.get_typed_unchecked(ctx, generator, idx, None))
Ok(v.get_typed_unchecked(ctx, generator, &idx, None))
}
},
).unwrap()
@ -803,7 +803,7 @@ pub fn ndarray_elementwise_binop_impl<'ctx, G, ValueFn>(
},
|generator, ctx, v, idx| {
unsafe {
Ok(v.dim_sizes().get_typed_unchecked(ctx, generator, idx, None))
Ok(v.dim_sizes().get_typed_unchecked(ctx, generator, &idx, None))
}
},
).unwrap()

View File

@ -189,7 +189,7 @@ pub fn gen_store_target<'ctx, G: CodeGenerator>(
[Some(raw_index), Some(len), None],
slice.location,
);
v.data().ptr_offset(ctx, generator, index, name)
v.data().ptr_offset(ctx, generator, &index, name)
}
TypeEnum::TObj { obj_id, .. } if *obj_id == PRIMITIVE_DEF_IDS.ndarray => {

View File

@ -1730,7 +1730,7 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
arg.dim_sizes().get_typed_unchecked(
ctx,
generator,
llvm_usize.const_zero(),
&llvm_usize.const_zero(),
None,
)
};