use crate::ast::expression::HamelinExpression;
use hamelin_lib::err::{TranslationError, TranslationErrors, UnexpectedType};
use hamelin_lib::sql::expression::apply::FunctionCallApply;
use hamelin_lib::sql::expression::literal::StringLiteral;
use hamelin_lib::sql::expression::Leaf;
use hamelin_lib::sql::expression::SQLExpression;
use hamelin_lib::types::array::Array;
use hamelin_lib::types::map::Map;
use hamelin_lib::types::{Type, UNKNOWN, VARIANT};
use hamelin_sql::utils::hamelin_array_index_to_sql_with_negatives;
use thiserror::Error;
pub struct HamelinIndexAccess {
pub value: HamelinExpression,
pub index: HamelinExpression,
}
impl HamelinIndexAccess {
pub fn new(value: HamelinExpression, index: HamelinExpression) -> Self {
Self { value, index }
}
pub fn infer_type(&self) -> Result<Type, TranslationErrors> {
match self.value.translate()?.typ {
Type::Array(Array { element_type }) => Ok(*element_type),
Type::Map(Map {
key_type: _key_type,
value_type,
}) => Ok(*value_type),
Type::Variant => Ok(VARIANT),
t => TranslationError::wrap(
self.value.tree(),
UnexpectedType::new(
t,
vec![
Array::new(UNKNOWN).into(),
Map::new(UNKNOWN, UNKNOWN).into(),
VARIANT,
],
),
)
.single_result(),
}
}
pub fn to_sql(&self) -> Result<SQLExpression, TranslationErrors> {
let value = self.value.translate()?;
let value_type = value.typ;
let value_sql = value.sql;
let index = self.index.translate()?;
let index_type = index.typ;
let index_sql = index.sql;
let rendered_index_sql = index_sql.to_string();
let expected_types = vec![
Array::new(UNKNOWN).into(),
Map::new(UNKNOWN, UNKNOWN).into(),
VARIANT,
];
match (value_type, index_type) {
(Type::Array(_), Type::Int) => Ok(FunctionCallApply::with_two(
"element_at",
value_sql.clone(),
hamelin_array_index_to_sql_with_negatives(index_sql, value_sql),
)
.into()),
(Type::Map(map), index_type) if *(map.key_type) == index_type => {
Ok(FunctionCallApply::with_two("element_at", value_sql, index_sql).into())
}
(Type::Variant, Type::Int) => {
if let SQLExpression::FunctionCallApply(json_extract) = value_sql.clone() {
if json_extract.function_name == "json_extract"
&& json_extract.arguments.len() == 2
{
let path = json_extract.arguments.last().unwrap().clone();
if let SQLExpression::Leaf(Leaf::StringLiteral(StringLiteral { value })) =
path
{
return Ok(FunctionCallApply::with_two(
&json_extract.function_name.clone(),
json_extract.arguments.first().unwrap().clone(),
StringLiteral::new(&format!("{}[{}]", value, rendered_index_sql))
.into(),
)
.into());
}
}
}
return Ok(FunctionCallApply::with_two(
"json_extract",
value_sql,
StringLiteral::new(&format!("$[{}]", rendered_index_sql)).into(),
)
.into());
}
(Type::Array(_), t) => {
let inner =
BadIndexAccess::new(t, "Index must be an integer for array index access.");
TranslationError::wrap(self.index.tree(), inner).single_result()
}
(
Type::Map(Map {
key_type,
value_type: _value_type,
}),
_,
) => {
let inner = BadIndexAccess::new(
*_value_type,
&format!(
"Index must match key type {} for map index access.",
key_type
),
);
TranslationError::wrap(self.index.tree(), inner).single_result()
}
(Type::Variant, t) => {
let inner =
BadIndexAccess::new(t, "Index must be integer for variant index access.");
TranslationError::wrap(self.index.tree(), inner).single_result()
}
(vt, _) => {
let inner = UnexpectedType::new(vt, expected_types);
TranslationError::wrap(self.value.tree(), inner).single_result()
}
}
}
}
#[derive(Error, Debug)]
#[error("Cannot perform index access with {typ}: {reason}")]
pub struct BadIndexAccess {
pub typ: Type,
pub reason: String,
}
impl BadIndexAccess {
pub fn new(typ: Type, reason: &str) -> Self {
Self {
typ: typ,
reason: reason.to_string(),
}
}
}