Skip to main content

kcl_lib/std/
array.rs

1use indexmap::IndexMap;
2
3use crate::ExecutorContext;
4use crate::NodePath;
5use crate::SourceRange;
6use crate::errors::KclError;
7use crate::errors::KclErrorDetails;
8use crate::execution::ControlFlowKind;
9use crate::execution::ExecState;
10use crate::execution::fn_call::Arg;
11use crate::execution::fn_call::Args;
12use crate::execution::kcl_value::FunctionSource;
13use crate::execution::kcl_value::KclValue;
14use crate::execution::types::RuntimeType;
15
16/// Apply a function to each element of an array.
17pub async fn map(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
18    let array: Vec<KclValue> = args.get_unlabeled_kw_arg("array", &RuntimeType::any_array(), exec_state)?;
19    let f: FunctionSource = args.get_kw_arg("f", &RuntimeType::function(), exec_state)?;
20    let new_array = inner_map(array, f, exec_state, &args).await?;
21    Ok(KclValue::HomArray {
22        value: new_array,
23        ty: RuntimeType::any(),
24    })
25}
26
27async fn inner_map(
28    array: Vec<KclValue>,
29    f: FunctionSource,
30    exec_state: &mut ExecState,
31    args: &Args,
32) -> Result<Vec<KclValue>, KclError> {
33    let mut new_array = Vec::with_capacity(array.len());
34    for elem in array {
35        let new_elem = call_map_closure(
36            elem,
37            &f,
38            args.source_range,
39            args.node_path.clone(),
40            exec_state,
41            &args.ctx,
42        )
43        .await?;
44        new_array.push(new_elem);
45    }
46    Ok(new_array)
47}
48
49async fn call_map_closure(
50    input: KclValue,
51    map_fn: &FunctionSource,
52    source_range: SourceRange,
53    node_path: Option<NodePath>,
54    exec_state: &mut ExecState,
55    ctxt: &ExecutorContext,
56) -> Result<KclValue, KclError> {
57    let args = Args::new(
58        Default::default(),
59        vec![(None, Arg::new(input, source_range))],
60        source_range,
61        node_path,
62        exec_state,
63        ctxt.clone(),
64        Some("map closure".to_owned()),
65    );
66    let output = map_fn.call_kw(None, exec_state, ctxt, args, source_range).await?;
67    let source_ranges = vec![source_range];
68    let output = output.ok_or_else(|| {
69        KclError::new_semantic(KclErrorDetails::new(
70            "Map function must return a value".to_owned(),
71            source_ranges,
72        ))
73    })?;
74    let output = match output.control {
75        ControlFlowKind::Continue => output.into_value(),
76        ControlFlowKind::Exit => {
77            let message = "Early return inside map function is currently not supported".to_owned();
78            debug_assert!(false, "{}", &message);
79            return Err(KclError::new_internal(KclErrorDetails::new(
80                message,
81                vec![source_range],
82            )));
83        }
84    };
85    Ok(output)
86}
87
88/// For each item in an array, update a value.
89pub async fn reduce(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
90    let array: Vec<KclValue> = args.get_unlabeled_kw_arg("array", &RuntimeType::any_array(), exec_state)?;
91    let f: FunctionSource = args.get_kw_arg("f", &RuntimeType::function(), exec_state)?;
92    let initial: KclValue = args.get_kw_arg("initial", &RuntimeType::any(), exec_state)?;
93    inner_reduce(array, initial, f, exec_state, &args).await
94}
95
96async fn inner_reduce(
97    array: Vec<KclValue>,
98    initial: KclValue,
99    f: FunctionSource,
100    exec_state: &mut ExecState,
101    args: &Args,
102) -> Result<KclValue, KclError> {
103    let mut reduced = initial;
104    for elem in array {
105        reduced = call_reduce_closure(
106            elem,
107            reduced,
108            &f,
109            args.source_range,
110            args.node_path.clone(),
111            exec_state,
112            &args.ctx,
113        )
114        .await?;
115    }
116
117    Ok(reduced)
118}
119
120async fn call_reduce_closure(
121    elem: KclValue,
122    accum: KclValue,
123    reduce_fn: &FunctionSource,
124    source_range: SourceRange,
125    node_path: Option<NodePath>,
126    exec_state: &mut ExecState,
127    ctxt: &ExecutorContext,
128) -> Result<KclValue, KclError> {
129    // Call the reduce fn for this repetition.
130    let mut labeled = IndexMap::with_capacity(1);
131    labeled.insert("accum".to_string(), Arg::new(accum, source_range));
132    let reduce_fn_args = Args::new(
133        labeled,
134        vec![(None, Arg::new(elem, source_range))],
135        source_range,
136        node_path,
137        exec_state,
138        ctxt.clone(),
139        Some("reduce closure".to_owned()),
140    );
141    let transform_fn_return = reduce_fn
142        .call_kw(None, exec_state, ctxt, reduce_fn_args, source_range)
143        .await?;
144
145    // Unpack the returned transform object.
146    let source_ranges = vec![source_range];
147    let out = transform_fn_return.ok_or_else(|| {
148        KclError::new_semantic(KclErrorDetails::new(
149            "Reducer function must return a value".to_string(),
150            source_ranges.clone(),
151        ))
152    })?;
153    let out = match out.control {
154        ControlFlowKind::Continue => out.into_value(),
155        ControlFlowKind::Exit => {
156            let message = "Early return inside reduce function is currently not supported".to_owned();
157            debug_assert!(false, "{}", &message);
158            return Err(KclError::new_internal(KclErrorDetails::new(
159                message,
160                vec![source_range],
161            )));
162        }
163    };
164    Ok(out)
165}
166
167pub async fn push(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
168    let (mut array, ty) = args.get_unlabeled_kw_arg_array_and_type("array", exec_state)?;
169    let item: KclValue = args.get_kw_arg("item", &RuntimeType::any(), exec_state)?;
170
171    array.push(item);
172
173    Ok(KclValue::HomArray { value: array, ty })
174}
175
176pub async fn pop(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
177    let (mut array, ty) = args.get_unlabeled_kw_arg_array_and_type("array", exec_state)?;
178    if array.is_empty() {
179        return Err(KclError::new_semantic(KclErrorDetails::new(
180            "Cannot pop from an empty array".to_string(),
181            vec![args.source_range],
182        )));
183    }
184    array.pop();
185    Ok(KclValue::HomArray { value: array, ty })
186}
187
188pub async fn concat(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
189    let (left, left_el_ty) = args.get_unlabeled_kw_arg_array_and_type("array", exec_state)?;
190    let right_value: KclValue = args.get_kw_arg("items", &RuntimeType::any_array(), exec_state)?;
191
192    match right_value {
193        KclValue::HomArray {
194            value: right,
195            ty: right_el_ty,
196            ..
197        } => Ok(inner_concat(&left, &left_el_ty, &right, &right_el_ty)),
198        KclValue::Tuple { value: right, .. } => {
199            // Tuples are treated as arrays for concatenation.
200            Ok(inner_concat(&left, &left_el_ty, &right, &RuntimeType::any()))
201        }
202        // Any single value is a subtype of an array, so we can treat it as a
203        // single-element array.
204        _ => Ok(inner_concat(&left, &left_el_ty, &[right_value], &RuntimeType::any())),
205    }
206}
207
208fn inner_concat(
209    left: &[KclValue],
210    left_el_ty: &RuntimeType,
211    right: &[KclValue],
212    right_el_ty: &RuntimeType,
213) -> KclValue {
214    if left.is_empty() {
215        return KclValue::HomArray {
216            value: right.to_vec(),
217            ty: right_el_ty.clone(),
218        };
219    }
220    if right.is_empty() {
221        return KclValue::HomArray {
222            value: left.to_vec(),
223            ty: left_el_ty.clone(),
224        };
225    }
226    let mut new = left.to_vec();
227    new.extend_from_slice(right);
228    // Propagate the element type if we can.
229    let ty = if right_el_ty.subtype(left_el_ty) {
230        left_el_ty.clone()
231    } else if left_el_ty.subtype(right_el_ty) {
232        right_el_ty.clone()
233    } else {
234        RuntimeType::any()
235    };
236    KclValue::HomArray { value: new, ty }
237}
238
239pub async fn slice(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
240    let (array, ty) = args.get_unlabeled_kw_arg_array_and_type("array", exec_state)?;
241    let start: Option<i64> = args.get_kw_arg_opt("start", &RuntimeType::count(), exec_state)?;
242    let end: Option<i64> = args.get_kw_arg_opt("end", &RuntimeType::count(), exec_state)?;
243
244    if start.is_none() && end.is_none() {
245        return Err(KclError::new_semantic(KclErrorDetails::new(
246            "Either `start` or `end` must be provided".to_owned(),
247            vec![args.source_range],
248        )));
249    }
250
251    let Ok(len) = i64::try_from(array.len()) else {
252        return Err(KclError::new_semantic(KclErrorDetails::new(
253            format!("Array length {} exceeds maximum supported length", array.len()),
254            vec![args.source_range],
255        )));
256    };
257    let mut computed_start = start.unwrap_or(0);
258    let mut computed_end = end.unwrap_or(len);
259
260    // Negative indices count from the end.
261    if computed_start < 0 {
262        computed_start += len;
263    }
264    if computed_end < 0 {
265        computed_end += len;
266    }
267
268    fn empty_slice(ty: RuntimeType) -> KclValue {
269        KclValue::HomArray { value: Vec::new(), ty }
270    }
271
272    if computed_start < 0 {
273        computed_start = 0;
274    }
275    if computed_start >= len {
276        return Ok(empty_slice(ty));
277    }
278    if computed_end > len {
279        computed_end = len;
280    }
281    if computed_end < 0 {
282        return Ok(empty_slice(ty));
283    }
284
285    if computed_start >= computed_end {
286        return Ok(empty_slice(ty));
287    }
288
289    let Some(sliced) = array.get(computed_start as usize..computed_end as usize) else {
290        let message = "Failed to compute array slice".to_owned();
291        debug_assert!(false, "{message}");
292        return Err(KclError::new_internal(KclErrorDetails::new(
293            message,
294            vec![args.source_range],
295        )));
296    };
297    Ok(KclValue::HomArray {
298        value: sliced.to_vec(),
299        ty,
300    })
301}
302
303pub async fn flatten(exec_state: &mut ExecState, args: Args) -> Result<KclValue, KclError> {
304    let array_value: KclValue = args.get_unlabeled_kw_arg("array", &RuntimeType::any_array(), exec_state)?;
305    let mut flattened = Vec::new();
306
307    let (array, original_ty) = match array_value {
308        KclValue::HomArray { value, ty, .. } => (value, ty),
309        KclValue::Tuple { value, .. } => (value, RuntimeType::any()),
310        _ => (vec![array_value], RuntimeType::any()),
311    };
312    for elem in array {
313        match elem {
314            KclValue::HomArray { value, .. } => flattened.extend(value),
315            KclValue::Tuple { value, .. } => flattened.extend(value),
316            _ => flattened.push(elem),
317        }
318    }
319
320    let ty = infer_flattened_type(original_ty, &flattened);
321    Ok(KclValue::HomArray { value: flattened, ty })
322}
323
324/// Infer the type of a flattened array based on the original type and the
325/// types of the flattened values. Currently, we preserve the original type only
326/// if all flattened values have the same type as the original element type.
327/// Otherwise, we fall back to `any`.
328fn infer_flattened_type(original_ty: RuntimeType, values: &[KclValue]) -> RuntimeType {
329    for value in values {
330        if !value.has_type(&original_ty) {
331            return RuntimeType::any();
332        };
333    }
334
335    original_ty
336}