fix: change list_access_to_uplc to properly handle list discards

crates/aiken-lang/src/gen_uplc/builder.rs

@@ -1,7 +1,7 @@
 use std::{collections::HashMap, ops::Deref, rc::Rc};
 
 use indexmap::{IndexMap, IndexSet};
-use itertools::Itertools;
+use itertools::{Itertools, Position};
 use uplc::{
     ast::{Constant as UplcConstant, Name, Term, Type as UplcType},
     builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER},
@@ -1530,127 +1530,152 @@ pub fn list_access_to_uplc(
     error_term: Term<Name>,
 ) -> Term<Name> {
     let names_len = names_types_ids.len();
+    // Should never be expect level none on a list
+    assert!(!(matches!(expect_level, ExpectLevel::None) && is_list_accessor));
 
     let mut no_tailing_discards = names_types_ids
         .iter()
         .rev()
-        .skip_while(|(name, _, _)| name == "_")
+        .with_position()
+        .skip_while(|pos| match pos {
+            // Items are reversed order
+            Position::Last((name, _, _)) | Position::Middle((name, _, _)) => {
+                name == "_" && matches!(expect_level, ExpectLevel::None)
+            }
+            Position::First((name, _, _)) | Position::Only((name, _, _)) => {
+                name == "_" && (tail_present || matches!(expect_level, ExpectLevel::None))
+            }
+        })
+        .map(|position| match position {
+            Position::First(a) | Position::Middle(a) | Position::Last(a) | Position::Only(a) => a,
+        })
         .collect_vec();
 
-    // If the the is just discards and check_last_item then we check for empty list
+    // If is just discards and check_last_item then we check for empty list
     if no_tailing_discards.is_empty() {
-        if !tail_present && matches!(expect_level, ExpectLevel::Full | ExpectLevel::Items) {
+        if tail_present || matches!(expect_level, ExpectLevel::None) {
+            return term.lambda("_");
+        } else {
             return Term::var("empty_list")
                 .delayed_choose_list(term, error_term)
                 .lambda("empty_list");
-        } else {
-            return term.lambda("_");
         }
     }
 
     // reverse back to original order
     no_tailing_discards.reverse();
 
-    let no_tailing_len = no_tailing_discards.len();
-
     // If we cut off at least one element then that was tail and possibly some heads
     let tail_wasnt_cutoff = tail_present && no_tailing_discards.len() == names_len;
 
-    no_tailing_discards.into_iter().enumerate().rev().fold(
+    let tail_name = |id| format!("tail_id_{}", id);
-        term,
-        |acc, (index, (name, tipo, id))| {
-            let tail_name = format!("tail_index_{}_{}", index, id);
 
-            let head_list =
+    let head_item = |name, tipo: &Rc<Type>, tail_name: &str| {
-                if matches!(tipo.get_uplc_type(), UplcType::Pair(_, _)) && is_list_accessor {
+        if name == "_" {
-                    Term::head_list().apply(Term::var(tail_name.to_string()))
+            Term::unit()
-                } else if matches!(expect_level, ExpectLevel::Full) && error_term != Term::Error {
+        } else if matches!(tipo.get_uplc_type(), UplcType::Pair(_, _)) && is_list_accessor {
-                    convert_data_to_type_debug(
+            Term::head_list().apply(Term::var(tail_name.to_string()))
-                        Term::head_list().apply(Term::var(tail_name.to_string())),
+        } else if matches!(expect_level, ExpectLevel::Full) && error_term != Term::Error {
-                        &tipo.to_owned(),
+            convert_data_to_type_debug(
-                        error_term.clone(),
+                Term::head_list().apply(Term::var(tail_name.to_string())),
-                    )
+                &tipo.to_owned(),
-                } else {
+                error_term.clone(),
-                    convert_data_to_type(
+            )
-                        Term::head_list().apply(Term::var(tail_name.to_string())),
+        } else {
-                        &tipo.to_owned(),
+            convert_data_to_type(
-                    )
+                Term::head_list().apply(Term::var(tail_name.to_string())),
-                };
+                &tipo.to_owned(),
+            )
+        }
+    };
 
-            // handle tail case
+    // Remember we reverse here so the First or Only is the last item
-            // name is guaranteed to not be discard at this point
+    no_tailing_discards
-            if index == no_tailing_len - 1 && tail_wasnt_cutoff {
+        .into_iter()
-                // simply lambda for tail name
+        .rev()
-                acc.lambda(name)
+        .with_position()
-            } else if index == no_tailing_len - 1 {
+        .fold(term, |acc, position| {
-                // case for no tail
+            match position {
-                // name is guaranteed to not be discard at this point
+                Position::First((name, _, _)) | Position::Only((name, _, _))
+                    if tail_wasnt_cutoff =>
+                {
+                    // case for tail as the last item
+                    acc.lambda(name)
+                }
 
-                match expect_level {
+                Position::First((name, tipo, id)) | Position::Only((name, tipo, id)) => {
-                    ExpectLevel::None => acc.lambda(name).apply(head_list).lambda(tail_name),
+                    // case for no tail, but last item
-                    ExpectLevel::Full | ExpectLevel::Items => {
+                    let tail_name = tail_name(id);
-                        if error_term == Term::Error && tail_present {
+
-                            acc.lambda(name).apply(head_list).lambda(tail_name)
+                    let head_item = head_item(name, tipo, &tail_name);
-                        } else if tail_present {
+
-                            // Custom error instead of trying to do head_list on a possibly empty list.
+                    match expect_level {
-                            Term::var(tail_name.to_string())
+                        ExpectLevel::None => acc.lambda(name).apply(head_item).lambda(tail_name),
-                                .delayed_choose_list(
+
-                                    error_term.clone(),
+                        ExpectLevel::Full | ExpectLevel::Items => {
-                                    acc.lambda(name).apply(head_list),
+                            if error_term == Term::Error && tail_present {
-                                )
+                                // No need to check last item if tail was present
-                                .lambda(tail_name)
+                                acc.lambda(name).apply(head_item).lambda(tail_name)
-                        } else if error_term == Term::Error {
+                            } else if tail_present {
-                            // Check head is last item in this list
+                                // Custom error instead of trying to do head_item on a possibly empty list.
-                            Term::tail_list()
+                                Term::var(tail_name.to_string())
-                                .apply(Term::var(tail_name.to_string()))
+                                    .delayed_choose_list(
-                                .delayed_choose_list(acc, error_term.clone())
+                                        error_term.clone(),
-                                .lambda(name)
+                                        acc.lambda(name).apply(head_item),
-                                .apply(head_list)
+                                    )
-                                .lambda(tail_name)
+                                    .lambda(tail_name)
-                        } else {
+                            } else if error_term == Term::Error {
-                            // Custom error if list is not empty after this head
+                                // Check head is last item in this list
-                            Term::var(tail_name.to_string())
+                                Term::tail_list()
-                                .delayed_choose_list(
+                                    .apply(Term::var(tail_name.to_string()))
-                                    error_term.clone(),
+                                    .delayed_choose_list(acc, error_term.clone())
-                                    Term::tail_list()
+                                    .lambda(name)
-                                        .apply(Term::var(tail_name.to_string()))
+                                    .apply(head_item)
-                                        .delayed_choose_list(acc, error_term.clone())
+                                    .lambda(tail_name)
-                                        .lambda(name)
+                            } else {
-                                        .apply(head_list),
+                                // Custom error if list is not empty after this head
-                                )
+                                Term::var(tail_name.to_string())
-                                .lambda(tail_name)
+                                    .delayed_choose_list(
+                                        error_term.clone(),
+                                        Term::tail_list()
+                                            .apply(Term::var(tail_name.to_string()))
+                                            .delayed_choose_list(acc, error_term.clone())
+                                            .lambda(name)
+                                            .apply(head_item),
+                                    )
+                                    .lambda(tail_name)
+                            }
                         }
                     }
                 }
-            } else if name == "_" {
+
-                if matches!(expect_level, ExpectLevel::None) || error_term == Term::Error {
+                Position::Middle((name, tipo, id)) | Position::Last((name, tipo, id)) => {
-                    acc.apply(Term::tail_list().apply(Term::var(tail_name.to_string())))
+                    // case for every item except the last item
-                        .lambda(tail_name)
+                    let tail_name = tail_name(id);
-                } else {
+
-                    Term::var(tail_name.to_string())
+                    let head_item = head_item(name, tipo, &tail_name);
-                        .delayed_choose_list(
+
-                            error_term.clone(),
+                    if matches!(expect_level, ExpectLevel::None) || error_term == Term::Error {
-                            acc.apply(Term::tail_list().apply(Term::var(tail_name.to_string()))),
-                        )
-                        .lambda(tail_name)
-                }
-            } else if matches!(expect_level, ExpectLevel::None) || error_term == Term::Error {
-                acc.apply(Term::tail_list().apply(Term::var(tail_name.to_string())))
-                    .lambda(name)
-                    .apply(head_list)
-                    .lambda(tail_name)
-            } else {
-                Term::var(tail_name.to_string())
-                    .delayed_choose_list(
-                        error_term.clone(),
                         acc.apply(Term::tail_list().apply(Term::var(tail_name.to_string())))
                             .lambda(name)
-                            .apply(head_list),
+                            .apply(head_item)
-                    )
+                            .lambda(tail_name)
-                    .lambda(tail_name)
+                    } else {
+                        // case for a custom error if the list is empty at this point
+                        Term::var(tail_name.to_string())
+                            .delayed_choose_list(
+                                error_term.clone(),
+                                acc.apply(
+                                    Term::tail_list().apply(Term::var(tail_name.to_string())),
+                                )
+                                .lambda(name)
+                                .apply(head_item),
+                            )
+                            .lambda(tail_name)
+                    }
+                }
             }
-        },
+        })
-    )
 }
 
 pub fn apply_builtin_forces(mut term: Term<Name>, force_count: u32) -> Term<Name> {

examples/acceptance_tests/070/aiken.lock

@@ -0,0 +1,16 @@
+# This file was generated by Aiken
+# You typically do not need to edit this file
+
+[[requirements]]
+name = "aiken-lang/stdlib"
+version = "main"
+source = "github"
+
+[[packages]]
+name = "aiken-lang/stdlib"
+version = "main"
+requirements = []
+source = "github"
+
+[etags]
+"aiken-lang/stdlib@main" = [{ secs_since_epoch = 1706674613, nanos_since_epoch = 871553000 }, "cf946239d3dd481ed41f20e56bf24910b5229ea35aa171a708edc2a47fc20a7b"]

examples/acceptance_tests/070/aiken.toml

@@ -0,0 +1,8 @@
+name = "aiken-lang/acceptance_test_070"
+version = '0.0.0'
+description = ''
+
+[[dependencies]]
+name = 'aiken-lang/stdlib'
+version = 'main'
+source = 'github'

examples/acceptance_tests/070/lib/tests.ak

@@ -0,0 +1,47 @@
+use aiken/list
+use aiken/transaction.{InlineDatum, Input, OutputReference, TransactionId}
+
+type OtherInput {
+  output_reference: OutputReference,
+  other: Data,
+}
+
+type MyDatum<a> {
+  Constructor1(a)
+  Constructor2
+}
+
+test discard_partitions() {
+  let all_inputs =
+    [
+      OtherInput(OutputReference(TransactionId(#"aabb"), 2), 3),
+      OtherInput(OutputReference(TransactionId(#"aabbcc"), 3), 3),
+    ]
+
+  let own_out_ref = OutputReference(TransactionId(#"aabb"), 2)
+
+  expect ([_], other_inputs) =
+    list.partition(
+      all_inputs,
+      fn(input) { input.output_reference == own_out_ref },
+    )
+
+  let inputs: List<Input> =
+    []
+
+  list.all(
+    inputs,
+    fn(input) {
+      expect dat: MyDatum<Int> =
+        when input.output.datum is {
+          InlineDatum(d) -> d
+          _ -> fail @"Not an inline datum"
+        }
+
+      when dat is {
+        Constructor1 { .. } -> True
+        _ -> False
+      }
+    },
+  )
+}