1use super::lookup_utils::cmp_for_lookup;
18use crate::args::{ArgSchema, CoercionPolicy, ShapeKind};
19use crate::function::Function;
20use crate::traits::{ArgumentHandle, CalcValue, FunctionContext};
21use formualizer_common::{ArgKind, ExcelError, ExcelErrorKind, LiteralValue};
22use formualizer_macros::func_caps;
23
24fn approx_match_ascending(slice: &[LiteralValue], needle: &LiteralValue) -> Option<usize> {
28 if slice.is_empty() {
29 return None;
30 }
31 let mut lo: usize = 0;
32 let mut hi: usize = slice.len();
33 while lo < hi {
34 let mid = (lo + hi) / 2;
35 match cmp_for_lookup(&slice[mid], needle) {
36 Some(c) if c > 0 => hi = mid,
37 Some(_) => lo = mid + 1,
38 None => hi = mid,
39 }
40 }
41 if lo == 0 { None } else { Some(lo - 1) }
42}
43
44#[derive(Debug)]
112pub struct LookupFn;
113impl Function for LookupFn {
124 fn name(&self) -> &'static str {
125 "LOOKUP"
126 }
127 fn min_args(&self) -> usize {
128 2
129 }
130 func_caps!(PURE, LOOKUP);
131 fn arg_schema(&self) -> &'static [ArgSchema] {
132 use once_cell::sync::Lazy;
133 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
134 vec![
135 ArgSchema {
137 kinds: smallvec::smallvec![ArgKind::Any],
138 required: true,
139 by_ref: false,
140 shape: ShapeKind::Scalar,
141 coercion: CoercionPolicy::None,
142 max: None,
143 repeating: None,
144 default: None,
145 },
146 ArgSchema {
148 kinds: smallvec::smallvec![ArgKind::Any],
149 required: true,
150 by_ref: false,
151 shape: ShapeKind::Range,
152 coercion: CoercionPolicy::None,
153 max: None,
154 repeating: None,
155 default: None,
156 },
157 ArgSchema {
159 kinds: smallvec::smallvec![ArgKind::Any],
160 required: false,
161 by_ref: false,
162 shape: ShapeKind::Range,
163 coercion: CoercionPolicy::None,
164 max: None,
165 repeating: None,
166 default: None,
167 },
168 ]
169 });
170 &SCHEMA
171 }
172
173 fn eval<'a, 'b, 'c>(
174 &self,
175 args: &'c [ArgumentHandle<'a, 'b>],
176 ctx: &dyn FunctionContext<'b>,
177 ) -> Result<CalcValue<'b>, ExcelError> {
178 if args.len() < 2 {
179 return Ok(CalcValue::Scalar(LiteralValue::Error(ExcelError::new(
180 ExcelErrorKind::Na,
181 ))));
182 }
183
184 let lookup_value = args[0].value()?.into_literal();
185 if let LiteralValue::Error(e) = lookup_value {
186 return Ok(CalcValue::Scalar(LiteralValue::Error(e)));
187 }
188
189 let has_result_vector = args.len() >= 3;
190
191 let lookup_data = materialise_range(&args[1], ctx)?;
193 let (l_rows, l_cols) = dims(&lookup_data);
194
195 let (search_vec, is_row_search) = if has_result_vector {
197 flatten_1d(&lookup_data, l_rows, l_cols)
199 } else if l_rows == 1 && l_cols == 1 {
200 (vec![lookup_data[0][0].clone()], false)
202 } else if l_cols > l_rows {
203 (lookup_data[0].clone(), true)
205 } else {
206 (
208 lookup_data.iter().map(|r| r[0].clone()).collect::<Vec<_>>(),
209 false,
210 )
211 };
212
213 let match_idx = approx_match_ascending(&search_vec, &lookup_value);
215 let match_idx = match match_idx {
216 Some(i) => i,
217 None => {
218 return Ok(CalcValue::Scalar(LiteralValue::Error(ExcelError::new(
219 ExcelErrorKind::Na,
220 ))));
221 }
222 };
223
224 if has_result_vector {
226 let result_data = materialise_range(&args[2], ctx)?;
227 let (r_rows, r_cols) = dims(&result_data);
228 let result_vec = flatten_1d_vec(&result_data, r_rows, r_cols);
229 let val = result_vec
230 .get(match_idx)
231 .cloned()
232 .unwrap_or(LiteralValue::Empty);
233 Ok(CalcValue::Scalar(materialise_empty(val)))
234 } else if l_rows == 1 && l_cols == 1 {
235 Ok(CalcValue::Scalar(materialise_empty(
236 lookup_data[0][0].clone(),
237 )))
238 } else if is_row_search {
239 let last_row = l_rows - 1;
241 let val = lookup_data
242 .get(last_row)
243 .and_then(|r| r.get(match_idx))
244 .cloned()
245 .unwrap_or(LiteralValue::Empty);
246 Ok(CalcValue::Scalar(materialise_empty(val)))
247 } else {
248 let last_col = l_cols - 1;
250 let val = lookup_data
251 .get(match_idx)
252 .and_then(|r| r.get(last_col))
253 .cloned()
254 .unwrap_or(LiteralValue::Empty);
255 Ok(CalcValue::Scalar(materialise_empty(val)))
256 }
257 }
258}
259
260fn materialise_range<'a, 'b>(
266 arg: &ArgumentHandle<'a, 'b>,
267 ctx: &dyn FunctionContext<'b>,
268) -> Result<Vec<Vec<LiteralValue>>, ExcelError> {
269 if let Ok(r) = arg.as_reference_or_eval() {
270 let current_sheet = ctx.current_sheet();
271 let rv = ctx.resolve_range_view(&r, current_sheet)?;
272 let (rows, cols) = rv.dims();
273 let mut data = Vec::with_capacity(rows);
274 rv.for_each_row(&mut |row| {
275 let mut owned = Vec::with_capacity(cols);
276 owned.extend_from_slice(row);
277 data.push(owned);
278 Ok(())
279 })?;
280 Ok(data)
281 } else {
282 let v = arg.value()?.into_literal();
283 match v {
284 LiteralValue::Array(rows) => Ok(rows),
285 other => Ok(vec![vec![other]]),
286 }
287 }
288}
289
290fn dims(data: &[Vec<LiteralValue>]) -> (usize, usize) {
291 let rows = data.len();
292 let cols = data.first().map(|r| r.len()).unwrap_or(0);
293 (rows, cols)
294}
295
296fn flatten_1d(data: &[Vec<LiteralValue>], rows: usize, cols: usize) -> (Vec<LiteralValue>, bool) {
299 if rows == 1 {
300 (data[0].clone(), true)
301 } else if cols == 1 {
302 (data.iter().map(|r| r[0].clone()).collect(), false)
303 } else {
304 (data.iter().flat_map(|r| r.iter().cloned()).collect(), false)
307 }
308}
309
310fn flatten_1d_vec(data: &[Vec<LiteralValue>], rows: usize, cols: usize) -> Vec<LiteralValue> {
312 flatten_1d(data, rows, cols).0
313}
314
315fn materialise_empty(v: LiteralValue) -> LiteralValue {
317 match v {
318 LiteralValue::Empty => LiteralValue::Number(0.0),
319 other => other,
320 }
321}
322
323#[cfg(test)]
328mod tests {
329 use super::*;
330 use crate::test_workbook::TestWorkbook;
331 use crate::traits::ArgumentHandle;
332 use formualizer_parse::parser::{ASTNode, ASTNodeType, ReferenceType};
333 use std::sync::Arc;
334
335 fn lit(v: LiteralValue) -> ASTNode {
336 ASTNode::new(ASTNodeType::Literal(v), None)
337 }
338
339 #[test]
342 fn approx_empty_slice() {
343 assert_eq!(approx_match_ascending(&[], &LiteralValue::Int(1)), None);
344 }
345
346 #[test]
347 fn approx_below_minimum() {
348 let vals = vec![
349 LiteralValue::Int(10),
350 LiteralValue::Int(20),
351 LiteralValue::Int(30),
352 ];
353 assert_eq!(approx_match_ascending(&vals, &LiteralValue::Int(5)), None);
354 }
355
356 #[test]
357 fn approx_exact_hit() {
358 let vals = vec![
359 LiteralValue::Int(10),
360 LiteralValue::Int(20),
361 LiteralValue::Int(30),
362 ];
363 assert_eq!(
364 approx_match_ascending(&vals, &LiteralValue::Int(20)),
365 Some(1)
366 );
367 }
368
369 #[test]
370 fn approx_between_values() {
371 let vals = vec![
372 LiteralValue::Int(10),
373 LiteralValue::Int(20),
374 LiteralValue::Int(30),
375 ];
376 assert_eq!(
377 approx_match_ascending(&vals, &LiteralValue::Int(25)),
378 Some(1)
379 );
380 }
381
382 #[test]
383 fn approx_above_max() {
384 let vals = vec![
385 LiteralValue::Int(10),
386 LiteralValue::Int(20),
387 LiteralValue::Int(30),
388 ];
389 assert_eq!(
390 approx_match_ascending(&vals, &LiteralValue::Int(100)),
391 Some(2)
392 );
393 }
394
395 #[test]
398 fn lookup_vector_exact_match() {
399 let wb = TestWorkbook::new()
400 .with_function(Arc::new(LookupFn))
401 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
402 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(2))
403 .with_cell_a1("Sheet1", "A3", LiteralValue::Int(3))
404 .with_cell_a1("Sheet1", "B1", LiteralValue::Text("a".into()))
405 .with_cell_a1("Sheet1", "B2", LiteralValue::Text("b".into()))
406 .with_cell_a1("Sheet1", "B3", LiteralValue::Text("c".into()));
407 let ctx = wb.interpreter();
408
409 let lookup_vec = ASTNode::new(
410 ASTNodeType::Reference {
411 original: "A1:A3".into(),
412 reference: ReferenceType::range(None, Some(1), Some(1), Some(3), Some(1)),
413 },
414 None,
415 );
416 let result_vec = ASTNode::new(
417 ASTNodeType::Reference {
418 original: "B1:B3".into(),
419 reference: ReferenceType::range(None, Some(1), Some(2), Some(3), Some(2)),
420 },
421 None,
422 );
423
424 let f = ctx.context.get_function("", "LOOKUP").unwrap();
425 let needle = lit(LiteralValue::Int(2));
426 let args = vec![
427 ArgumentHandle::new(&needle, &ctx),
428 ArgumentHandle::new(&lookup_vec, &ctx),
429 ArgumentHandle::new(&result_vec, &ctx),
430 ];
431 let v = f
432 .dispatch(&args, &ctx.function_context(None))
433 .unwrap()
434 .into_literal();
435 assert_eq!(v, LiteralValue::Text("b".into()));
436 }
437
438 #[test]
439 fn lookup_vector_approximate() {
440 let wb = TestWorkbook::new()
441 .with_function(Arc::new(LookupFn))
442 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
443 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(2))
444 .with_cell_a1("Sheet1", "A3", LiteralValue::Int(3))
445 .with_cell_a1("Sheet1", "A4", LiteralValue::Int(4))
446 .with_cell_a1("Sheet1", "A5", LiteralValue::Int(5));
447 let ctx = wb.interpreter();
448
449 let lookup_vec = ASTNode::new(
450 ASTNodeType::Reference {
451 original: "A1:A5".into(),
452 reference: ReferenceType::range(None, Some(1), Some(1), Some(5), Some(1)),
453 },
454 None,
455 );
456
457 let f = ctx.context.get_function("", "LOOKUP").unwrap();
458 let needle = lit(LiteralValue::Number(3.5));
459 let args = vec![
460 ArgumentHandle::new(&needle, &ctx),
461 ArgumentHandle::new(&lookup_vec, &ctx),
462 ];
463 let v = f
464 .dispatch(&args, &ctx.function_context(None))
465 .unwrap()
466 .into_literal();
467 assert_eq!(v, LiteralValue::Number(3.0));
469 }
470
471 #[test]
472 fn lookup_vector_below_min_returns_na() {
473 let wb = TestWorkbook::new()
474 .with_function(Arc::new(LookupFn))
475 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(10))
476 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(20));
477 let ctx = wb.interpreter();
478
479 let lookup_vec = ASTNode::new(
480 ASTNodeType::Reference {
481 original: "A1:A2".into(),
482 reference: ReferenceType::range(None, Some(1), Some(1), Some(2), Some(1)),
483 },
484 None,
485 );
486
487 let f = ctx.context.get_function("", "LOOKUP").unwrap();
488 let needle = lit(LiteralValue::Int(5));
489 let args = vec![
490 ArgumentHandle::new(&needle, &ctx),
491 ArgumentHandle::new(&lookup_vec, &ctx),
492 ];
493 let v = f
494 .dispatch(&args, &ctx.function_context(None))
495 .unwrap()
496 .into_literal();
497 assert!(matches!(v, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Na));
498 }
499
500 #[test]
503 fn lookup_array_column_search() {
504 let wb = TestWorkbook::new().with_function(Arc::new(LookupFn));
506 let ctx = wb.interpreter();
507
508 let arr = lit(LiteralValue::Array(vec![
509 vec![LiteralValue::Int(1), LiteralValue::Text("a".into())],
510 vec![LiteralValue::Int(2), LiteralValue::Text("b".into())],
511 vec![LiteralValue::Int(3), LiteralValue::Text("c".into())],
512 ]));
513
514 let f = ctx.context.get_function("", "LOOKUP").unwrap();
515 let needle = lit(LiteralValue::Int(2));
516 let args = vec![
517 ArgumentHandle::new(&needle, &ctx),
518 ArgumentHandle::new(&arr, &ctx),
519 ];
520 let v = f
521 .dispatch(&args, &ctx.function_context(None))
522 .unwrap()
523 .into_literal();
524 assert_eq!(v, LiteralValue::Text("b".into()));
525 }
526
527 #[test]
528 fn lookup_array_row_search() {
529 let wb = TestWorkbook::new().with_function(Arc::new(LookupFn));
531 let ctx = wb.interpreter();
532
533 let arr = lit(LiteralValue::Array(vec![
534 vec![
535 LiteralValue::Int(1),
536 LiteralValue::Int(2),
537 LiteralValue::Int(3),
538 ],
539 vec![
540 LiteralValue::Text("x".into()),
541 LiteralValue::Text("y".into()),
542 LiteralValue::Text("z".into()),
543 ],
544 ]));
545
546 let f = ctx.context.get_function("", "LOOKUP").unwrap();
547 let needle = lit(LiteralValue::Int(2));
548 let args = vec![
549 ArgumentHandle::new(&needle, &ctx),
550 ArgumentHandle::new(&arr, &ctx),
551 ];
552 let v = f
553 .dispatch(&args, &ctx.function_context(None))
554 .unwrap()
555 .into_literal();
556 assert_eq!(v, LiteralValue::Text("y".into()));
557 }
558
559 #[test]
560 fn lookup_error_propagation() {
561 let wb = TestWorkbook::new().with_function(Arc::new(LookupFn));
562 let ctx = wb.interpreter();
563
564 let arr = lit(LiteralValue::Array(vec![vec![
565 LiteralValue::Int(1),
566 LiteralValue::Int(2),
567 ]]));
568 let needle = lit(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
569 let f = ctx.context.get_function("", "LOOKUP").unwrap();
570 let args = vec![
571 ArgumentHandle::new(&needle, &ctx),
572 ArgumentHandle::new(&arr, &ctx),
573 ];
574 let v = f
575 .dispatch(&args, &ctx.function_context(None))
576 .unwrap()
577 .into_literal();
578 assert!(matches!(v, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
579 }
580
581 #[test]
582 fn lookup_single_element() {
583 let wb = TestWorkbook::new().with_function(Arc::new(LookupFn));
584 let ctx = wb.interpreter();
585
586 let arr = lit(LiteralValue::Array(vec![vec![LiteralValue::Int(5)]]));
587 let f = ctx.context.get_function("", "LOOKUP").unwrap();
588
589 let needle = lit(LiteralValue::Int(5));
591 let args = vec![
592 ArgumentHandle::new(&needle, &ctx),
593 ArgumentHandle::new(&arr, &ctx),
594 ];
595 let v = f
596 .dispatch(&args, &ctx.function_context(None))
597 .unwrap()
598 .into_literal();
599 assert_eq!(v, LiteralValue::Int(5));
600
601 let needle_lo = lit(LiteralValue::Int(3));
603 let args2 = vec![
604 ArgumentHandle::new(&needle_lo, &ctx),
605 ArgumentHandle::new(&arr, &ctx),
606 ];
607 let v2 = f
608 .dispatch(&args2, &ctx.function_context(None))
609 .unwrap()
610 .into_literal();
611 assert!(matches!(v2, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Na));
612 }
613
614 #[test]
615 fn lookup_text_values() {
616 let wb = TestWorkbook::new()
618 .with_function(Arc::new(LookupFn))
619 .with_cell_a1("Sheet1", "A1", LiteralValue::Text("apple".into()))
620 .with_cell_a1("Sheet1", "A2", LiteralValue::Text("banana".into()))
621 .with_cell_a1("Sheet1", "A3", LiteralValue::Text("cherry".into()))
622 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(1))
623 .with_cell_a1("Sheet1", "B2", LiteralValue::Int(2))
624 .with_cell_a1("Sheet1", "B3", LiteralValue::Int(3));
625 let ctx = wb.interpreter();
626
627 let lookup_vec = ASTNode::new(
628 ASTNodeType::Reference {
629 original: "A1:A3".into(),
630 reference: ReferenceType::range(None, Some(1), Some(1), Some(3), Some(1)),
631 },
632 None,
633 );
634 let result_vec = ASTNode::new(
635 ASTNodeType::Reference {
636 original: "B1:B3".into(),
637 reference: ReferenceType::range(None, Some(1), Some(2), Some(3), Some(2)),
638 },
639 None,
640 );
641
642 let f = ctx.context.get_function("", "LOOKUP").unwrap();
643 let needle = lit(LiteralValue::Text("banana".into()));
644 let args = vec![
645 ArgumentHandle::new(&needle, &ctx),
646 ArgumentHandle::new(&lookup_vec, &ctx),
647 ArgumentHandle::new(&result_vec, &ctx),
648 ];
649 let v = f
650 .dispatch(&args, &ctx.function_context(None))
651 .unwrap()
652 .into_literal();
653 assert_eq!(v, LiteralValue::Number(2.0));
654 }
655}