Skip to main content

openjd_expr/functions/
list.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5//! List function implementations (sorted, reversed, unique, flatten, range).
6
7use crate::error::ExpressionError;
8use crate::function_library::EvalContext;
9use crate::types::ExprType;
10use crate::value::ExprValue;
11
12type R = Result<ExprValue, ExpressionError>;
13type Ctx<'a> = &'a mut dyn EvalContext;
14
15pub fn sorted_fn(ctx: Ctx, a: &[ExprValue]) -> R {
16    let (elements, elem_type) = a[0]
17        .clone()
18        .into_list()
19        .ok_or_else(|| ExpressionError::new("sorted() argument must be a list"))?;
20    ctx.count_ops(elements.len())?;
21    let mut sorted = elements;
22    sorted.sort_by(|a, b| a.compare(b).unwrap_or(std::cmp::Ordering::Equal));
23    ExprValue::make_list_checked(ctx, sorted, elem_type)
24}
25
26pub fn reversed_fn(ctx: Ctx, a: &[ExprValue]) -> R {
27    let (mut elements, elem_type) = a[0]
28        .clone()
29        .into_list()
30        .ok_or_else(|| ExpressionError::new("reversed() argument must be a list"))?;
31    ctx.count_ops(elements.len())?;
32    elements.reverse();
33    ExprValue::make_list_checked(ctx, elements, elem_type)
34}
35
36pub fn unique_fn(ctx: Ctx, a: &[ExprValue]) -> R {
37    let (elements, elem_type) = a[0]
38        .clone()
39        .into_list()
40        .ok_or_else(|| ExpressionError::new("unique() argument must be a list"))?;
41    ctx.count_ops(elements.len())?;
42    let mut seen = std::collections::HashSet::new();
43    let mut result = Vec::new();
44    for e in elements {
45        if seen.insert(e.clone()) {
46            result.push(e);
47        }
48    }
49    ExprValue::make_list_checked(ctx, result, elem_type)
50}
51
52pub fn flatten_fn(ctx: Ctx, a: &[ExprValue]) -> R {
53    let iter = a[0]
54        .list_iter()
55        .ok_or_else(|| ExpressionError::new("flatten() argument must be a list"))?;
56    let mut result = Vec::new();
57    for e in iter {
58        ctx.count_op()?;
59        if e.is_list() {
60            let inner_len = e.list_len().unwrap_or(0);
61            ctx.count_ops(inner_len)?;
62            result.extend(e.list_iter().expect("is_list() was true"));
63        } else {
64            result.push(e);
65        }
66    }
67    let et = if result.is_empty() {
68        ExprType::INT
69    } else {
70        result[0].expr_type()
71    };
72    ExprValue::make_list_checked(ctx, result, et)
73}
74
75pub fn range_fn(ctx: Ctx, a: &[ExprValue]) -> R {
76    let (start, stop, step) = match a.len() {
77        1 => (
78            0i64,
79            match &a[0] {
80                ExprValue::Int(n) => *n,
81                _ => return Err(ExpressionError::new("range() argument must be int")),
82            },
83            1i64,
84        ),
85        2 => (
86            match &a[0] {
87                ExprValue::Int(n) => *n,
88                _ => return Err(ExpressionError::new("range() arguments must be int")),
89            },
90            match &a[1] {
91                ExprValue::Int(n) => *n,
92                _ => return Err(ExpressionError::new("range() arguments must be int")),
93            },
94            1,
95        ),
96        3 => (
97            match &a[0] {
98                ExprValue::Int(n) => *n,
99                _ => return Err(ExpressionError::new("range() arguments must be int")),
100            },
101            match &a[1] {
102                ExprValue::Int(n) => *n,
103                _ => return Err(ExpressionError::new("range() arguments must be int")),
104            },
105            match &a[2] {
106                ExprValue::Int(n) => *n,
107                _ => return Err(ExpressionError::new("range() arguments must be int")),
108            },
109        ),
110        _ => return Err(ExpressionError::new("range() takes 1-3 arguments")),
111    };
112    if step == 0 {
113        return Err(ExpressionError::new("range() step cannot be zero"));
114    }
115    let mut elements = Vec::new();
116    let mut v = start;
117    if step > 0 {
118        while v < stop {
119            elements.push(ExprValue::Int(v));
120            v += step;
121            ctx.count_op()?;
122        }
123    } else {
124        while v > stop {
125            elements.push(ExprValue::Int(v));
126            v += step;
127            ctx.count_op()?;
128        }
129    }
130    ExprValue::make_list_checked(ctx, elements, ExprType::INT)
131}
132
133pub fn join_fn(ctx: Ctx, a: &[ExprValue]) -> R {
134    let iter = a[0]
135        .list_iter()
136        .ok_or_else(|| ExpressionError::new("join() argument must be a list"))?;
137    let sep = match a.get(1) {
138        Some(ExprValue::String(s)) => s.as_str(),
139        _ => "",
140    };
141    let mut parts = Vec::new();
142    for e in iter {
143        ctx.count_op()?;
144        parts.push(e.to_display_string());
145    }
146    Ok(ExprValue::String(parts.join(sep)))
147}
148
149pub fn list_from_range(ctx: Ctx, a: &[ExprValue]) -> R {
150    match &a[0] {
151        ExprValue::RangeExpr(r) => {
152            ctx.count_ops(r.len())?;
153            let elements: Vec<ExprValue> = r.iter().map(ExprValue::Int).collect();
154            Ok(ExprValue::make_list_checked(ctx, elements, ExprType::INT)?)
155        }
156        _ => Err(ExpressionError::new("list() argument must be range_expr")),
157    }
158}
159
160pub fn range_expr_from_string(_: Ctx, a: &[ExprValue]) -> R {
161    match &a[0] {
162        ExprValue::String(s) => {
163            let r: crate::range_expr::RangeExpr = s.parse()?;
164            Ok(ExprValue::RangeExpr(r))
165        }
166        _ => Err(ExpressionError::new("range_expr() argument must be string")),
167    }
168}
169
170pub fn range_expr_from_list(ctx: Ctx, a: &[ExprValue]) -> R {
171    if let Some(iter) = a[0].list_iter() {
172        if a[0].list_len() == Some(0) {
173            return Err(ExpressionError::new(
174                "range_expr() requires at least one value",
175            ));
176        }
177        ctx.count_ops(a[0].list_len().unwrap_or(0))?;
178        let mut ints: Vec<i64> = Vec::new();
179        for e in iter {
180            match e {
181                ExprValue::Int(i) => ints.push(i),
182                _ => return Err(ExpressionError::new("range_expr() list must contain ints")),
183            }
184        }
185        ints.sort();
186        ints.dedup();
187        let r = crate::range_expr::RangeExpr::from_values(ints);
188        Ok(ExprValue::RangeExpr(r))
189    } else {
190        Err(ExpressionError::new("range_expr() argument must be list"))
191    }
192}
193
194pub fn range_expr_from_empty_list(_: Ctx, _a: &[ExprValue]) -> R {
195    Err(ExpressionError::new(
196        "range_expr() requires at least one value",
197    ))
198}