From f62cb13089802131a6a5527eb27262038c642244 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Mon, 12 Aug 2024 10:43:18 -0400
Subject: [PATCH] Allow indexing lists with ranges

---
 dust-lang/src/analyzer.rs | 18 +++++++++++-------
 dust-lang/src/value.rs    |  8 ++++++++
 dust-lang/src/vm.rs       | 31 ++++++++++++++++++++++++++-----
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/dust-lang/src/analyzer.rs b/dust-lang/src/analyzer.rs
index daf0e94..7771b83 100644
--- a/dust-lang/src/analyzer.rs
+++ b/dust-lang/src/analyzer.rs
@@ -334,10 +334,14 @@ impl<'a> Analyzer<'a> {
                 }
 
                 if let Some(Type::List { .. }) = left.inner.expected_type(self.context) {
-                    if let Some(Type::Integer) = right.inner.expected_type(self.context) {
+                    let right_type = right.inner.expected_type(self.context);
+
+                    if let Some(Type::Integer) = right_type {
                         // Allow indexing lists with integers
+                    } else if let Some(Type::Range) = right_type {
+                        // Allow indexing lists with ranges
                     } else {
-                        return Err(AnalyzerError::ExpectedInteger {
+                        return Err(AnalyzerError::ExpectedIntegerOrRange {
                             actual: right.as_ref().clone(),
                         });
                     }
@@ -408,7 +412,7 @@ pub enum AnalyzerError {
     ExpectedIdentifierOrString {
         actual: Node<Statement>,
     },
-    ExpectedInteger {
+    ExpectedIntegerOrRange {
         actual: Node<Statement>,
     },
     ExpectedValue {
@@ -441,7 +445,7 @@ impl AnalyzerError {
             AnalyzerError::ExpectedBoolean { actual, .. } => actual.position,
             AnalyzerError::ExpectedIdentifier { actual, .. } => actual.position,
             AnalyzerError::ExpectedIdentifierOrString { actual } => actual.position,
-            AnalyzerError::ExpectedInteger { actual, .. } => actual.position,
+            AnalyzerError::ExpectedIntegerOrRange { actual, .. } => actual.position,
             AnalyzerError::ExpectedValue { actual } => actual.position,
             AnalyzerError::ExpectedValueArgumentCount { position, .. } => *position,
             AnalyzerError::TypeConflict {
@@ -468,8 +472,8 @@ impl Display for AnalyzerError {
             AnalyzerError::ExpectedIdentifierOrString { actual } => {
                 write!(f, "Expected identifier or string, found {}", actual)
             }
-            AnalyzerError::ExpectedInteger { actual, .. } => {
-                write!(f, "Expected integer, found {}", actual)
+            AnalyzerError::ExpectedIntegerOrRange { actual, .. } => {
+                write!(f, "Expected integer or range, found {}", actual)
             }
             AnalyzerError::ExpectedValue { actual, .. } => {
                 write!(f, "Expected value, found {}", actual)
@@ -514,7 +518,7 @@ mod tests {
         assert_eq!(
             analyze(source),
             Err(DustError::AnalyzerError {
-                analyzer_error: AnalyzerError::ExpectedInteger {
+                analyzer_error: AnalyzerError::ExpectedIntegerOrRange {
                     actual: Node::new(Statement::Identifier(Identifier::new("foo")), (10, 13)),
                 },
                 source
diff --git a/dust-lang/src/value.rs b/dust-lang/src/value.rs
index f80c6bf..1d37a08 100644
--- a/dust-lang/src/value.rs
+++ b/dust-lang/src/value.rs
@@ -145,6 +145,14 @@ impl Value {
         }
     }
 
+    pub fn as_range(&self) -> Option<&Range<i64>> {
+        if let ValueInner::Range(range) = self.inner().as_ref() {
+            Some(range)
+        } else {
+            None
+        }
+    }
+
     pub fn add(&self, other: &Value) -> Result<Value, ValueError> {
         match (self.inner().as_ref(), other.inner().as_ref()) {
             (ValueInner::Float(left), ValueInner::Float(right)) => Ok(Value::float(left + right)),
diff --git a/dust-lang/src/vm.rs b/dust-lang/src/vm.rs
index df690ac..d8ac15a 100644
--- a/dust-lang/src/vm.rs
+++ b/dust-lang/src/vm.rs
@@ -485,6 +485,14 @@ impl Vm {
 
                         return Ok(value);
                     }
+
+                    if let Some(range) = value.as_range() {
+                        let range = range.start as usize..range.end as usize;
+
+                        if let Some(items) = list.get(range) {
+                            return Ok(Some(Value::list(items.to_vec())));
+                        }
+                    }
                 }
 
                 if let (Some(map), Statement::Identifier(identifier)) =
@@ -495,7 +503,7 @@ impl Vm {
                     return Ok(value);
                 }
 
-                Err(VmError::ExpectedIdentifierOrInteger {
+                Err(VmError::ExpectedIdentifierIntegerOrRange {
                     position: right_span,
                 })
             }
@@ -575,7 +583,7 @@ pub enum VmError {
     ExpectedIdentifier {
         position: Span,
     },
-    ExpectedIdentifierOrInteger {
+    ExpectedIdentifierIntegerOrRange {
         position: Span,
     },
     ExpectedInteger {
@@ -613,7 +621,7 @@ impl VmError {
             Self::BuiltInFunctionError { position, .. } => *position,
             Self::ExpectedBoolean { position } => *position,
             Self::ExpectedIdentifier { position } => *position,
-            Self::ExpectedIdentifierOrInteger { position } => *position,
+            Self::ExpectedIdentifierIntegerOrRange { position } => *position,
             Self::ExpectedInteger { position } => *position,
             Self::ExpectedFunction { position, .. } => *position,
             Self::ExpectedList { position } => *position,
@@ -654,10 +662,10 @@ impl Display for VmError {
             Self::ExpectedIdentifier { position } => {
                 write!(f, "Expected an identifier at position: {:?}", position)
             }
-            Self::ExpectedIdentifierOrInteger { position } => {
+            Self::ExpectedIdentifierIntegerOrRange { position } => {
                 write!(
                     f,
-                    "Expected an identifier or integer at position: {:?}",
+                    "Expected an identifier, integer, or range at position: {:?}",
                     position
                 )
             }
@@ -693,6 +701,19 @@ impl Display for VmError {
 mod tests {
     use super::*;
 
+    #[test]
+    fn list_index_range() {
+        let input = "[1, 2, 3, 4, 5].1..3";
+
+        assert_eq!(
+            run(input),
+            Ok(Some(Value::list(vec![
+                Value::integer(2),
+                Value::integer(3)
+            ])))
+        );
+    }
+
     #[test]
     fn range() {
         let input = "1..5";