2 files changed, 64 insertions, 13 deletions

diff --git a/pkg/shell/interpreter.go b/pkg/shell/interpreter.go
index dc0faf9..07375a2 100644
--- a/pkg/shell/interpreter.go
+++ b/pkg/shell/interpreter.go
@@ -12,6 +12,7 @@ import (
 	"unicode"
 
 	"github.com/jpappel/atlas/pkg/data"
+	"github.com/jpappel/atlas/pkg/index"
 	"github.com/jpappel/atlas/pkg/query"
 	"github.com/jpappel/atlas/pkg/util"
 	"golang.org/x/term"
@@ -533,16 +534,18 @@ out:
 				return true, errors.New("Type corruption during compilation, expected query.CompilationArtifact")
 			}
 
-			results, err := inter.querier.Execute(context.Background(), artifact)
+			resultsMap, err := inter.querier.Execute(context.Background(), artifact)
 			if err != nil {
 				return false, fmt.Errorf("Error occured while excuting query: %s", err)
 			}
 
-			_, err = query.YamlOutput{}.OutputTo(w, slices.Collect(maps.Values(results)))
-			if err != nil {
-				return false, fmt.Errorf("Can't output results: %s", err)
-			}
-			fmt.Fprintln(w)
+			results := slices.Collect(maps.Values(resultsMap))
+			stack = append(stack, Value{VAL_RESULTS, results})
+			// _, err = query.YamlOutput{}.OutputTo(w, slices.Collect(maps.Values(results)))
+			// if err != nil {
+			// 	return false, fmt.Errorf("Can't output results: %s", err)
+			// }
+			// fmt.Fprintln(w)
 		case ITOK_VAR_NAME:
 			val, ok := inter.State[t.Text]
 			if !ok {
@@ -594,6 +597,12 @@ out:
 					return true, fmt.Errorf("Type corruption during len, expected *query.Clause")
 				}
 				length = clause.Order()
+			case VAL_RESULTS:
+				results, ok := arg.Val.([]*index.Document)
+				if !ok {
+					return true, fmt.Errorf("Type corruption during len, expected []*index.Document")
+				}
+				length = len(results)
 			default:
 				return false, fmt.Errorf("Unable to get length of argument with type %s", arg.Type)
 			}
@@ -646,6 +655,14 @@ out:
 						pos++
 					}
 				}
+			case VAL_RESULTS:
+				if results, ok := arg.Val.([]*index.Document); !ok {
+					return true, fmt.Errorf("Type corruption during at, expecetd []*index.Document")
+				} else if idx < 0 || idx >= len(results) {
+					return false, fmt.Errorf("Index out of bounds")
+				} else {
+					stack = append(stack, Value{VAL_RESULTS, results[idx : idx+1]})
+				}
 			default:
 				return false, fmt.Errorf("Cannot index type %s", arg.Type)
 			}
@@ -704,6 +721,20 @@ out:
 						startIdx, stopIdx, len(qTokens),
 					)
 				}
+			case VAL_RESULTS:
+				results, ok := arg.Val.([]*index.Document)
+				if !ok {
+					return true, fmt.Errorf("Type corruption during slice, expected []*index.Document")
+				}
+
+				if 0 <= startIdx && startIdx <= stopIdx && stopIdx <= len(results) {
+					stack = append(stack, Value{VAL_RESULTS, results[startIdx:stopIdx]})
+				} else {
+					return false, fmt.Errorf(
+						"Indexes [%d:%d] out of range [0:%d]",
+						startIdx, stopIdx, len(results),
+					)
+				}
 			default:
 				return false, fmt.Errorf("Cannot slice argument of type %s", arg.Type)
 			}
diff --git a/pkg/shell/state.go b/pkg/shell/state.go
index 5acdaf1..1eed62f 100644
--- a/pkg/shell/state.go
+++ b/pkg/shell/state.go
@@ -1,9 +1,11 @@
 package shell
 
 import (
+	"bytes"
 	"fmt"
 	"strings"
 
+	"github.com/jpappel/atlas/pkg/index"
 	"github.com/jpappel/atlas/pkg/query"
 )
 
@@ -16,6 +18,7 @@ const (
 	VAL_TOKENS
 	VAL_CLAUSE
 	VAL_ARTIFACT
+	VAL_RESULTS
 )
 
 type Value struct {
@@ -39,6 +42,8 @@ func (t ValueType) String() string {
 		return "Clause"
 	case VAL_ARTIFACT:
 		return "Compilation Artifact"
+	case VAL_RESULTS:
+		return "Query Result"
 	default:
 		return "Unkown"
 	}
@@ -76,6 +81,18 @@ func (v Value) String() string {
 			panic("Corrupted Type (expected query.CompilationArtifact)")
 		}
 		return artifact.String()
+	case VAL_RESULTS:
+		results, ok := v.Val.([]*index.Document)
+		if !ok {
+			panic("Corrupted Type (expected []*index.Document)")
+		}
+
+		b := bytes.Buffer{}
+		yo := query.YamlOutput{}
+		if _, err := yo.OutputTo(&b, results); err != nil {
+			panic(err)
+		}
+		return b.String()
 	case VAL_INVALID:
 		return "Invalid"
 	}
@@ -88,21 +105,24 @@ func (s State) String() string {
 	for k, v := range s {
 		b.WriteString(k)
 		b.WriteByte(':')
+		b.WriteByte(' ')
 		switch v.Type {
 		case VAL_INVALID:
-			b.WriteString(" Invalid")
+			b.WriteString("Invalid")
 		case VAL_INT:
-			b.WriteString(" Integer")
+			b.WriteString("Integer")
 		case VAL_STRING:
-			b.WriteString(" String")
+			b.WriteString("String")
 		case VAL_TOKENS:
-			b.WriteString(" Tokens")
+			b.WriteString("Tokens")
 		case VAL_CLAUSE:
-			b.WriteString(" Clause")
+			b.WriteString("Clause")
 		case VAL_ARTIFACT:
-			b.WriteString(" Artifact")
+			b.WriteString("Artifact")
+		case VAL_RESULTS:
+			b.WriteString("Results")
 		default:
-			fmt.Fprintf(&b, " Unknown (%d)", v.Val)
+			fmt.Fprintf(&b, "Unknown (%d)", v.Val)
 		}
 		b.WriteByte('\n')
 	}