Add eject host function (#231)

eigerco · Jan 20, 2025 · 6dcf292 · 6dcf292
1 parent ce23ec9
commit 6dcf292
Show file tree

Hide file tree

Showing 2 changed files with 124 additions and 2 deletions.
diff --git a/internal/polkavm/host_call/accumulate_functions.go b/internal/polkavm/host_call/accumulate_functions.go
@@ -1,13 +1,16 @@
 package host_call
 
 import (
+	"bytes"
+
 	"github.com/eigerco/strawberry/internal/block"
 	"github.com/eigerco/strawberry/internal/common"
 	"github.com/eigerco/strawberry/internal/crypto"
 	"github.com/eigerco/strawberry/internal/jamtime"
 	. "github.com/eigerco/strawberry/internal/polkavm"
 	"github.com/eigerco/strawberry/internal/service"
 	"github.com/eigerco/strawberry/internal/state"
+	"github.com/eigerco/strawberry/pkg/serialization/codec/jam"
 )
 
 // Bless ΩB(ϱ, ω, μ, (x, y))
@@ -253,9 +256,63 @@ func Eject(gas Gas, regs Registers, mem Memory, ctxPair AccumulateContextPair, t
 	}
 	gas -= EjectCost
 
-	// TODO: implement method
+	d, o := regs[A0], regs[A1]
 
-	return gas, regs, mem, ctxPair, nil
+	// let h = μo..o+32 if Zo..o+32 ⊂ Vμ
+	h := make([]byte, 32)
+	if err := mem.Read(uint32(o), h); err != nil {
+		// otherwise ∇
+		return gas, withCode(regs, OOB), mem, ctxPair, nil
+	}
+
+	if block.ServiceId(d) == ctxPair.RegularCtx.ServiceId {
+		// d = x_s => WHO
+		return gas, withCode(regs, WHO), mem, ctxPair, nil
+	}
+
+	// if d ∈ K((x_u)_d)
+	serviceAccount, ok := ctxPair.RegularCtx.AccumulationState.ServiceState[block.ServiceId(d)]
+	if !ok {
+		return gas, withCode(regs, WHO), mem, ctxPair, nil
+	}
+
+	encodedXs, err := jam.Marshal(struct {
+		ServiceId block.ServiceId `jam:"length=32"`
+	}{ctxPair.RegularCtx.ServiceId})
+	if err != nil || !bytes.Equal(serviceAccount.CodeHash[:], encodedXs) {
+		// d_c ≠ E32(x_s) => WHO
+		return gas, withCode(regs, WHO), mem, ctxPair, err
+	}
+
+	if serviceAccount.TotalItems() != 2 {
+		// d_i ≠ 2 => HUH
+		return gas, withCode(regs, HUH), mem, ctxPair, nil
+	}
+
+	// l = max(81, d_o) - 81
+	l := max(81, len(serviceAccount.Code())) - 81
+
+	key := service.PreImageMetaKey{Hash: crypto.Hash(h), Length: service.PreimageLength(l)}
+	dL, ok := serviceAccount.PreimageMeta[key]
+	if !ok {
+		// (h, l) ∉ d_l => HUH
+		return gas, withCode(regs, HUH), mem, ctxPair, nil
+	}
+
+	// if d_l[h, l] = [x, y], y < t − D => OK
+	if len(dL) == 2 && dL[1] < timeslot-jamtime.PreimageExpulsionPeriod {
+		xs := ctxPair.RegularCtx.ServiceAccount()
+		// s'_b = ((x_u)d)[x_s]b + d_b
+		xs.Balance += serviceAccount.Balance
+
+		delete(ctxPair.RegularCtx.AccumulationState.ServiceState, block.ServiceId(d))
+		ctxPair.RegularCtx.AccumulationState.ServiceState[ctxPair.RegularCtx.ServiceId] = xs
+
+		return gas, withCode(regs, OK), mem, ctxPair, nil
+	}
+
+	// otherwise => HUH
+	return gas, withCode(regs, HUH), mem, ctxPair, nil
 }
 
 // Query ΩQ(ϱ, ω, μ, (x, y))

diff --git a/internal/polkavm/host_call/accumulate_functions_test.go b/internal/polkavm/host_call/accumulate_functions_test.go
@@ -282,6 +282,71 @@ func TestAccumulate(t *testing.T) {
 				}},
 			},
 		},
+		{
+			name:       "eject",
+			fn:         fnTms(Eject),
+			initialGas: 100,
+			timeslot:   200,
+			initialRegs: deltaRegs{
+				A0: 999,
+			},
+			alloc: alloc{
+				A1: hash2bytes(randomHash),
+			},
+			X: AccumulateContext{
+				ServiceId: 222,
+				AccumulationState: state.AccumulationState{
+					ServiceState: service.ServiceState{
+						// d = 999 (the ejected service)
+						999: func() service.ServiceAccount {
+							e32, err := jam.Marshal(struct {
+								ServiceId block.ServiceId `jam:"length=32"`
+							}{222})
+							require.NoError(t, err)
+
+							var codeHash crypto.Hash
+							copy(codeHash[:], e32)
+
+							preImgLookup := map[crypto.Hash][]byte{
+								codeHash: make([]byte, 81),
+							}
+
+							preImgMeta := map[service.PreImageMetaKey]service.PreimageHistoricalTimeslots{
+								{Hash: randomHash, Length: 0}: {50, 100},
+							}
+
+							return service.ServiceAccount{
+								CodeHash:       codeHash,
+								Balance:        100, // d_b
+								PreimageLookup: preImgLookup,
+								PreimageMeta:   preImgMeta,
+							}
+						}(),
+						// x_s
+						222: {
+							Balance: 1000,
+						},
+					},
+				},
+			},
+			expectedGas: 88,
+			expectedDeltaRegs: deltaRegs{
+				A0: uint64(OK),
+			},
+			// After success:
+			//   - The ejected service 999 is removed
+			//   - x_s(222).Balance = 1000+100=1100
+			expectedX: AccumulateContext{
+				ServiceId: 222,
+				AccumulationState: state.AccumulationState{
+					ServiceState: service.ServiceState{
+						222: {
+							Balance: 1100,
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "solicit_out_of_gas",
 			fn:   fnTms(Solicit),