occheung/nac3
1
0
Fork 0
forked from M-Labs/nac3
Code Issues Pull Requests Projects Releases Wiki Activity

core: Implement type inference for indexing into ndarray

Browse Source
This commit is contained in:
David Mak 2024-02-19 17:07:18 +08:00
parent 976a9512c1
commit 0d5c53e60c
1 changed files with 90 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
nac3core/src/typecheck/type_inferencer/mod.rs
View File

@ -1237,6 +1237,67 @@ impl<'a> Inferencer<'a> {
Ok(boolean)
}
/// Infers the type of a subscript expression on an `ndarray`.
fn infer_subscript_ndarray(
&mut self,
value: &ast::Expr<Option<Type>>,
dummy_tvar: Type,
ndims: &Type,
) -> InferenceResult {
debug_assert!(matches!(
&*self.unifier.get_ty_immutable(dummy_tvar),
TypeEnum::TVar { is_const_generic: false, .. }
));
let constrained_ty = self.unifier.add_ty(TypeEnum::TNDArray { ty: dummy_tvar, ndims: *ndims });
self.constrain(value.custom.unwrap(), constrained_ty, &value.location)?;
let TypeEnum::TLiteral { values, .. } = &*self.unifier.get_ty_immutable(*ndims) else {
panic!("Expected TLiteral for TNDArray.ndims, got {}", self.unifier.stringify(*ndims))
};
let ndims = values.iter()
.map(|ndim| match *ndim {
SymbolValue::U64(v) => Ok(v),
SymbolValue::U32(v) => Ok(v as u64),
SymbolValue::I32(v) => u64::try_from(v).map_err(|_| HashSet::from([
format!("Expected non-negative literal for TNDArray.ndims, got {v}"),
])),
SymbolValue::I64(v) => u64::try_from(v).map_err(|_| HashSet::from([
format!("Expected non-negative literal for TNDArray.ndims, got {v}"),
])),
_ => unreachable!(),
})
.collect::<Result<Vec<_>, _>>()?;
assert!(!ndims.is_empty());
if ndims.len() == 1 && ndims[0] == 1 {
// ndarray[T, Literal[1]] - Index always returns an object of type T
assert_ne!(ndims[0], 0);
Ok(dummy_tvar)
} else {
// ndarray[T, Literal[N]] where N != 1 - Index returns an object of type ndarray[T, Literal[N - 1]]
if ndims.iter().any(|v| *v == 0) {
unimplemented!("Inference for ndarray subscript operator with Literal[0, ...] bound unimplemented")
}
let ndims_min_one_ty = self.unifier.get_fresh_literal(
ndims.into_iter().map(|v| SymbolValue::U64(v - 1)).collect(),
None,
);
let subscripted_ty = self.unifier.add_ty(TypeEnum::TNDArray {
ty: dummy_tvar,
ndims: ndims_min_one_ty,
});
Ok(subscripted_ty)
}
}
fn infer_subscript(
&mut self,
value: &ast::Expr<Option<Type>>,
@ -1258,6 +1319,9 @@ impl<'a> Inferencer<'a> {
Ok(list_like_ty)
}
ExprKind::Constant { value: ast::Constant::Int(val), .. } => {
if let TypeEnum::TNDArray { ndims, .. } = &*self.unifier.get_ty(value.custom.unwrap()) {
self.infer_subscript_ndarray(value, ty, ndims)
} else {
// the index is a constant, so value can be a sequence.
let ind: Option<i32> = (*val).try_into().ok();
let ind = ind.ok_or_else(|| HashSet::from(["Index must be int32".to_string()]))?;
@ -1270,22 +1334,27 @@ impl<'a> Inferencer<'a> {
self.constrain(value.custom.unwrap(), seq, &value.location)?;
Ok(ty)
}
}
_ => {
if let TypeEnum::TTuple { .. } = &*self.unifier.get_ty(value.custom.unwrap())
{
if let TypeEnum::TTuple { .. } = &*self.unifier.get_ty(value.custom.unwrap()) {
return report_error("Tuple index must be a constant (KernelInvariant is also not supported)", slice.location)
}
// the index is not a constant, so value can only be a list
// the index is not a constant, so value can only be a list-like structure
match &*self.unifier.get_ty(value.custom.unwrap()) {
TypeEnum::TList { .. } => {
self.constrain(slice.custom.unwrap(), self.primitives.int32, &slice.location)?;
let list_like_ty = match &*self.unifier.get_ty(value.custom.unwrap()) {
TypeEnum::TList { .. } => self.unifier.add_ty(TypeEnum::TList { ty }),
TypeEnum::TNDArray { .. } => todo!(),
_ => unreachable!(),
};
self.constrain(value.custom.unwrap(), list_like_ty, &value.location)?;
let list = self.unifier.add_ty(TypeEnum::TList { ty });
self.constrain(value.custom.unwrap(), list, &value.location)?;
Ok(ty)
}
TypeEnum::TNDArray { ndims, .. } => {
self.constrain(slice.custom.unwrap(), self.primitives.usize(), &slice.location)?;
self.infer_subscript_ndarray(value, ty, ndims)
}
_ => unreachable!(),
}
}
}
}