fix: remove key from tb.keys and tb.k2i

Author: reversearrow

table.go

@@ -208,19 +208,36 @@ func (tb *LTable) RawSetString(key string, value LValue) {
 		tb.k2i = map[LValue]int{}
 	}
 
+	lkey := LString(key)
 	if value == LNil {
-		// TODO tb.keys and tb.k2i should also be removed
 		delete(tb.strdict, key)
+		tb.deleteKeyFromKeys(lkey)
+		delete(tb.k2i, lkey)
 	} else {
 		tb.strdict[key] = value
-		lkey := LString(key)
 		if _, ok := tb.k2i[lkey]; !ok {
 			tb.k2i[lkey] = len(tb.keys)
 			tb.keys = append(tb.keys, lkey)
 		}
 	}
 }
 
+func (tb *LTable) deleteKeyFromKeys(key LValue) {
+	idx, ok := tb.k2i[key]
+	if !ok {
+		return
+	}
+
+	lastIdx := len(tb.keys) - 1
+	if idx != lastIdx {
+		lastKey := tb.keys[lastIdx]
+		tb.keys[idx] = lastKey
+		tb.k2i[lastKey] = idx
+	}
+
+	tb.keys = tb.keys[:lastIdx]
+}
+
 // RawSetH sets a given LValue to a given index without the __newindex metamethod.
 func (tb *LTable) RawSetH(key LValue, value LValue) {
 	if s, ok := key.(LString); ok {
@@ -236,8 +253,9 @@ func (tb *LTable) RawSetH(key LValue, value LValue) {
 	}
 
 	if value == LNil {
-		// TODO tb.keys and tb.k2i should also be removed
 		delete(tb.dict, key)
+		tb.deleteKeyFromKeys(key)
+		delete(tb.k2i, key)
 	} else {
 		tb.dict[key] = value
 		if _, ok := tb.k2i[key]; !ok {

table_test.go

@@ -231,3 +231,249 @@ func TestTableForEach(t *testing.T) {
 		}
 	})
 }
+
+func TestLTable_RawSetString(t *testing.T) {
+	type input struct {
+		key   string
+		value LValue
+	}
+
+	tests := []struct {
+		name                  string
+		input                 []input
+		strdictExpectedLength int
+		expectedKeys          []LValue
+	}{
+		{
+			name: "input received as expected",
+			input: []input{
+				{
+					key:   "a",
+					value: LString("a"),
+				},
+				{
+					key:   "b",
+					value: LString("b"),
+				},
+				{
+					key:   "c",
+					value: LString("c"),
+				},
+				{
+					key:   "d",
+					value: LString("d"),
+				},
+			},
+			strdictExpectedLength: 4,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("c"), LString("d"),
+			},
+		},
+		{
+			name: "key removed from the middle",
+			input: []input{
+				{
+					key:   "a",
+					value: LString("a"),
+				},
+				{
+					key:   "b",
+					value: LString("b"),
+				},
+				{
+					key:   "c",
+					value: LString("c"),
+				},
+				{
+					key:   "d",
+					value: LString("d"),
+				},
+				{
+					key:   "c",
+					value: LNil,
+				},
+			},
+			strdictExpectedLength: 3,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("d"),
+			},
+		},
+		{
+			name: "key removed from the end",
+			input: []input{
+				{
+					key:   "a",
+					value: LString("a"),
+				},
+				{
+					key:   "b",
+					value: LString("b"),
+				},
+				{
+					key:   "c",
+					value: LString("c"),
+				},
+				{
+					key:   "d",
+					value: LString("d"),
+				},
+				{
+					key:   "d",
+					value: LNil,
+				},
+			},
+			strdictExpectedLength: 3,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("c"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tb := newLTable(0, 0)
+			for _, v := range tt.input {
+				tb.RawSetString(v.key, v.value)
+			}
+			if len(tb.strdict) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.strdict))
+			}
+
+			if len(tb.k2i) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.k2i))
+			}
+
+			if len(tb.keys) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.keys))
+			}
+
+			for i, v := range tb.keys {
+				if v != tt.expectedKeys[i] {
+					t.Error("Expected", tt.expectedKeys[i], "got", v)
+				}
+			}
+		})
+	}
+}
+
+func TestLTable_RawGetH(t *testing.T) {
+	type input struct {
+		key   LValue
+		value LValue
+	}
+
+	tests := []struct {
+		name                  string
+		input                 []input
+		strdictExpectedLength int
+		expectedKeys          []LValue
+	}{
+		{
+			name: "input received as expected",
+			input: []input{
+				{
+					key:   LString("a"),
+					value: LString("a"),
+				},
+				{
+					key:   LString("b"),
+					value: LString("b"),
+				},
+				{
+					key:   LString("c"),
+					value: LString("c"),
+				},
+				{
+					key:   LString("d"),
+					value: LString("d"),
+				},
+			},
+			strdictExpectedLength: 4,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("c"), LString("d"),
+			},
+		},
+		{
+			name: "key removed from the middle, expect removal from keys and k2i",
+			input: []input{
+				{
+					key:   LString("a"),
+					value: LString("a"),
+				},
+				{
+					key:   LString("b"),
+					value: LString("b"),
+				},
+				{
+					key:   LString("c"),
+					value: LString("c"),
+				},
+				{
+					key:   LString("d"),
+					value: LString("d"),
+				},
+				{
+					key:   LString("c"),
+					value: LNil,
+				},
+			},
+			strdictExpectedLength: 3,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("d"),
+			},
+		},
+		{
+			name: "key removed from the end, expect removal from keys and k2i",
+			input: []input{
+				{
+					key:   LString("a"),
+					value: LString("a"),
+				},
+				{
+					key:   LString("b"),
+					value: LString("b"),
+				},
+				{
+					key:   LString("c"),
+					value: LString("c"),
+				},
+				{
+					key:   LString("d"),
+					value: LString("d"),
+				},
+				{
+					key:   LString("d"),
+					value: LNil,
+				},
+			},
+			strdictExpectedLength: 3,
+			expectedKeys: []LValue{
+				LString("a"), LString("b"), LString("c"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tb := newLTable(0, 0)
+			for _, v := range tt.input {
+				tb.RawSetH(v.key, v.value)
+			}
+			if len(tb.strdict) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.strdict))
+			}
+
+			if len(tb.k2i) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.k2i))
+			}
+
+			if len(tb.keys) != tt.strdictExpectedLength {
+				t.Error("Expected length", tt.strdictExpectedLength, "got", len(tb.keys))
+			}
+
+			for i, v := range tb.keys {
+				if v != tt.expectedKeys[i] {
+					t.Error("Expected", tt.expectedKeys[i], "got", v)
+				}
+			}
+		})
+	}
+}