From b49404d9b640644f1b1e3ba42e113298d72d5a0e Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 26 Nov 2021 15:31:36 +0300
Subject: [PATCH] [#80] netmap: add `NodeInfo.ExitSubnet` method

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 netmap/subnet.go      | 11 ++++++++++-
 netmap/subnet_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/netmap/subnet.go b/netmap/subnet.go
index 09b7a57f..b61a504e 100644
--- a/netmap/subnet.go
+++ b/netmap/subnet.go
@@ -11,6 +11,15 @@ import (
 // EnterSubnet writes to NodeInfo the intention to enter the subnet. Must not be called on nil.
 // Zero NodeInfo belongs to zero subnet.
 func (i *NodeInfo) EnterSubnet(id subnetid.ID) {
+	i.changeSubnet(id, true)
+}
+
+// ExitSubnet writes to NodeInfo the intention to exit subnet. Must not be called on nil.
+func (i *NodeInfo) ExitSubnet(id subnetid.ID) {
+	i.changeSubnet(id, false)
+}
+
+func (i *NodeInfo) changeSubnet(id subnetid.ID, isMember bool) {
 	var (
 		idv2 refs.SubnetID
 		info netmap.NodeSubnetInfo
@@ -19,7 +28,7 @@ func (i *NodeInfo) EnterSubnet(id subnetid.ID) {
 	id.WriteToV2(&idv2)
 
 	info.SetID(&idv2)
-	info.SetEntryFlag(true)
+	info.SetEntryFlag(isMember)
 
 	netmap.WriteSubnetInfo((*netmap.NodeInfo)(i), info)
 }
diff --git a/netmap/subnet_test.go b/netmap/subnet_test.go
index cb1c0c47..0dd3600d 100644
--- a/netmap/subnet_test.go
+++ b/netmap/subnet_test.go
@@ -86,6 +86,30 @@ func TestNodeInfoSubnets(t *testing.T) {
 	})
 }
 
+func TestEnterSubnet(t *testing.T) {
+	var (
+		id   subnetid.ID
+		node netmap.NodeInfo
+	)
+
+	require.True(t, netmap.BelongsToSubnet(&node, id))
+
+	node.EnterSubnet(id)
+	require.True(t, netmap.BelongsToSubnet(&node, id))
+
+	node.ExitSubnet(id)
+	require.False(t, netmap.BelongsToSubnet(&node, id))
+
+	id.SetNumber(10)
+	node.EnterSubnet(id)
+	require.True(t, netmap.BelongsToSubnet(&node, id))
+	require.False(t, netmap.BelongsToSubnet(&node, subnetid.ID{}))
+
+	node.ExitSubnet(id)
+	require.False(t, netmap.BelongsToSubnet(&node, id))
+	require.False(t, netmap.BelongsToSubnet(&node, subnetid.ID{}))
+}
+
 func TestBelongsToSubnet(t *testing.T) {
 	var id, idMiss, idZero subnetid.ID