newell/nac3
1
0
Fork 0
forked from M-Labs/nac3
Code Pull Requests Activity

core: Add multidimensional array helpers

Browse Source
This commit is contained in:
David Mak 2024-06-04 18:00:42 +08:00
parent 2d4329e23c
commit 82bf14785b
1 changed files with 33 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
nac3core/src/toplevel/helper.rs
View File

@ -1,6 +1,7 @@
use std::convert::TryInto;
use crate::symbol_resolver::SymbolValue;
use crate::toplevel::numpy::unpack_ndarray_var_tys;
use crate::typecheck::typedef::{Mapping, VarMap};
use nac3parser::ast::{Constant, Location};
@ -691,3 +692,35 @@ pub fn parse_parameter_default_value(
]))
}
}
/// Obtains the element type of an array-like type.
pub fn arraylike_flatten_element_type(unifier: &mut Unifier, ty: Type) -> Type {
match &*unifier.get_ty(ty) {
TypeEnum::TObj { obj_id, .. } if *obj_id == PRIMITIVE_DEF_IDS.ndarray =>
unpack_ndarray_var_tys(unifier, ty).0,
TypeEnum::TList { ty } => arraylike_flatten_element_type(unifier, *ty),
_ => ty
}
}
/// Obtains the number of dimensions of an array-like type.
pub fn arraylike_get_ndims(unifier: &mut Unifier, ty: Type) -> u64 {
match &*unifier.get_ty(ty) {
TypeEnum::TObj { obj_id, .. } if *obj_id == PRIMITIVE_DEF_IDS.ndarray => {
let ndims = unpack_ndarray_var_tys(unifier, ty).1;
let TypeEnum::TLiteral { values, .. } = &*unifier.get_ty_immutable(ndims) else {
panic!("Expected TLiteral for ndarray.ndims, got {}", unifier.stringify(ndims))
};
if values.len() > 1 {
todo!("Getting num of dimensions for ndarray with more than one ndim bound is unimplemented")
}
u64::try_from(values[0].clone()).unwrap()
}
TypeEnum::TList { ty } => arraylike_get_ndims(unifier, *ty) + 1,
_ => 0
}
}