Optimize Kernel Invariant on Unwrap and Tuple #261

Merged
sb10q merged 1 commits from optimize_unwrap_kernelinvariant into master 2024-08-17 17:37:19 +08:00
3 changed files with 46 additions and 9 deletions

View File

@ -212,6 +212,28 @@ impl StaticValue for PythonValue {
})) }))
}) })
} }
fn get_tuple_element<'ctx>(&self, index: u32) -> Option<ValueEnum<'ctx>> {
Python::with_gil(|py| -> PyResult<Option<(u64, PyObject)>> {
let helper = &self.resolver.helper;
let ty = helper.type_fn.call1(py, (&self.value,))?;
let ty_id: u64 = helper.id_fn.call1(py, (ty,))?.extract(py)?;
assert_eq!(ty_id, self.resolver.primitive_ids.tuple);
let tup: &PyTuple = self.value.extract(py)?;
let elem = tup.get_item(index as usize);
let id = self.resolver.helper.id_fn.call1(py, (elem,))?.extract(py)?;
Ok(Some((id, elem.into())))
})
.unwrap()
.map(|(id, obj)| {
ValueEnum::Static(Arc::new(PythonValue {
id,
value: obj,
store_obj: self.store_obj.clone(),
resolver: self.resolver.clone(),
}))
})
}
} }
impl InnerResolver { impl InnerResolver {

View File

@ -1515,26 +1515,39 @@ pub fn gen_expr<'ctx, 'a, G: CodeGenerator>(
[Some(raw_index), Some(len), None], [Some(raw_index), Some(len), None],
expr.location, expr.location,
); );
ctx.build_gep_and_load(arr_ptr, &[index]) ctx.build_gep_and_load(arr_ptr, &[index]).into()
} }
} else if let TypeEnum::TTuple { .. } = &*ctx.unifier.get_ty(value.custom.unwrap()) { } else if let TypeEnum::TTuple { .. } = &*ctx.unifier.get_ty(value.custom.unwrap()) {
let v = generator
.gen_expr(ctx, value)?
.unwrap()
.to_basic_value_enum(ctx, generator, value.custom.unwrap())?
.into_struct_value();
let index: u32 = let index: u32 =
if let ExprKind::Constant { value: ast::Constant::Int(v), .. } = &slice.node { if let ExprKind::Constant { value: ast::Constant::Int(v), .. } = &slice.node {
(*v).try_into().unwrap() (*v).try_into().unwrap()
} else { } else {
unreachable!("tuple subscript must be const int after type check"); unreachable!("tuple subscript must be const int after type check");
}; };
ctx.builder.build_extract_value(v, index, "tup_elem").unwrap() let v = generator
.gen_expr(ctx, value)?
.unwrap();
match v {
ValueEnum::Dynamic(v) => {
let v = v.into_struct_value();
ctx.builder.build_extract_value(v, index, "tup_elem").unwrap().into()
}
ValueEnum::Static(v) => {
match v.get_tuple_element(index) {
Some(v) => v,
None => {
let tup = v
.to_basic_value_enum(ctx, generator, value.custom.unwrap())?
.into_struct_value();
ctx.builder.build_extract_value(tup, index, "tup_elem").unwrap().into()
}
}
}
}
} else { } else {
unreachable!("should not be other subscriptable types after type check"); unreachable!("should not be other subscriptable types after type check");
} }
} },
.into(),
ExprKind::ListComp { .. } => gen_comprehension(generator, ctx, expr)?.into(), ExprKind::ListComp { .. } => gen_comprehension(generator, ctx, expr)?.into(),
_ => unimplemented!(), _ => unimplemented!(),
})) }))

View File

@ -79,6 +79,8 @@ pub trait StaticValue {
name: StrRef, name: StrRef,
ctx: &mut CodeGenContext<'ctx, 'a>, ctx: &mut CodeGenContext<'ctx, 'a>,
) -> Option<ValueEnum<'ctx>>; ) -> Option<ValueEnum<'ctx>>;
fn get_tuple_element<'ctx>(&self, index: u32) -> Option<ValueEnum<'ctx>>;
} }
#[derive(Clone)] #[derive(Clone)]