From b792feccab1f8b107d329d83bf638e23cbb886e4 Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Thu, 18 Mar 2021 10:02:22 -0700
Subject: [PATCH] functions for indexing slices and strings

---
 src/Expr.cc | 107 +++++++++++++++++++++++++++++-----------------------
 src/Expr.h  |   4 ++
 2 files changed, 63 insertions(+), 48 deletions(-)

diff --git a/src/Expr.cc b/src/Expr.cc
index 1f5afbf8a7..24cfde992e 100644
--- a/src/Expr.cc
+++ b/src/Expr.cc
@@ -2936,24 +2936,7 @@ ValPtr IndexExpr::Fold(Val* v1, Val* v2) const
 		if ( lv->Length() == 1 )
 			v = vect->ValAt(lv->Idx(0)->CoerceToUnsigned());
 		else
-			{
-			size_t len = vect->Size();
-			auto result = make_intrusive<VectorVal>(vect->GetType<VectorType>());
-
-			bro_int_t first = get_slice_index(lv->Idx(0)->CoerceToInt(), len);
-			bro_int_t last = get_slice_index(lv->Idx(1)->CoerceToInt(), len);
-			bro_int_t sub_length = last - first;
-
-			if ( sub_length >= 0 )
-				{
-				result->Resize(sub_length);
-
-				for ( bro_int_t idx = first; idx < last; idx++ )
-					result->Assign(idx - first, vect->ValAt(idx));
-				}
-
-			return result;
-			}
+			return index_slice(vect, lv);
 		}
 		break;
 
@@ -2962,36 +2945,7 @@ ValPtr IndexExpr::Fold(Val* v1, Val* v2) const
 		break;
 
 	case TYPE_STRING:
-		{
-		const ListVal* lv = v2->AsListVal();
-		const String* s = v1->AsString();
-		int len = s->Len();
-		String* substring = nullptr;
-
-		if ( lv->Length() == 1 )
-			{
-			bro_int_t idx = lv->Idx(0)->AsInt();
-
-			if ( idx < 0 )
-				idx += len;
-
-			// Out-of-range index will return null pointer.
-			substring = s->GetSubstring(idx, 1);
-			}
-		else
-			{
-			bro_int_t first = get_slice_index(lv->Idx(0)->AsInt(), len);
-			bro_int_t last = get_slice_index(lv->Idx(1)->AsInt(), len);
-			bro_int_t substring_len = last - first;
-
-			if ( substring_len < 0 )
-				substring = nullptr;
-			else
-				substring = s->GetSubstring(first, substring_len);
-			}
-
-		return make_intrusive<StringVal>(substring ? substring : new String(""));
-		}
+		return index_string(v1->AsString(), v2->AsListVal());
 
 	default:
 		RuntimeError("type cannot be indexed");
@@ -3005,6 +2959,63 @@ ValPtr IndexExpr::Fold(Val* v1, Val* v2) const
 	return nullptr;
 	}
 
+StringValPtr index_string(const String* s, const ListVal* lv)
+	{
+	int len = s->Len();
+	String* substring = nullptr;
+
+	if ( lv->Length() == 1 )
+		{
+		bro_int_t idx = lv->Idx(0)->AsInt();
+
+		if ( idx < 0 )
+			idx += len;
+
+		// Out-of-range index will return null pointer.
+		substring = s->GetSubstring(idx, 1);
+		}
+	else
+		{
+		bro_int_t first = get_slice_index(lv->Idx(0)->AsInt(), len);
+		bro_int_t last = get_slice_index(lv->Idx(1)->AsInt(), len);
+		bro_int_t substring_len = last - first;
+
+		if ( substring_len < 0 )
+			substring = nullptr;
+		else
+			substring = s->GetSubstring(first, substring_len);
+		}
+
+	return make_intrusive<StringVal>(substring ? substring : new String(""));
+	}
+
+VectorValPtr index_slice(VectorVal* vect, const ListVal* lv)
+	{
+	auto first = lv->Idx(0)->CoerceToInt();
+	auto last = lv->Idx(1)->CoerceToInt();
+	return index_slice(vect, first, last);
+	}
+
+VectorValPtr index_slice(VectorVal* vect, int _first, int _last)
+	{
+	size_t len = vect->Size();
+	auto result = make_intrusive<VectorVal>(vect->GetType<VectorType>());
+
+	bro_int_t first = get_slice_index(_first, len);
+	bro_int_t last = get_slice_index(_last, len);
+	bro_int_t sub_length = last - first;
+
+	if ( sub_length >= 0 )
+		{
+		result->Resize(sub_length);
+
+		for ( bro_int_t idx = first; idx < last; idx++ )
+			result->Assign(idx - first, vect->ValAt(idx));
+		}
+
+	return result;
+	}
+
 void IndexExpr::Assign(Frame* f, ValPtr v)
 	{
 	if ( IsError() )
diff --git a/src/Expr.h b/src/Expr.h
index ae275a7a84..6b99b1572a 100644
--- a/src/Expr.h
+++ b/src/Expr.h
@@ -941,6 +941,10 @@ protected:
 	bool is_slice;
 };
 
+extern VectorValPtr index_slice(VectorVal* vect, const ListVal* lv);
+extern VectorValPtr index_slice(VectorVal* vect, int first, int last);
+extern StringValPtr index_string(const String* s, const ListVal* lv);
+
 class IndexExprWhen final : public IndexExpr {
 public:
 	static inline std::vector<ValPtr> results = {};