Commit 9cebc05a authored by Brian Tiger Chow's avatar Brian Tiger Chow

fix(bitswap/message) impl with map to ensure no duplicate blocks

comes at the cost of O(n) Blocks() method.
parent 824a185f
...@@ -26,15 +26,15 @@ type Exportable interface { ...@@ -26,15 +26,15 @@ type Exportable interface {
ToNet(p peer.Peer) (nm.NetMessage, error) ToNet(p peer.Peer) (nm.NetMessage, error)
} }
// message wraps a proto message for convenience
type impl struct { type impl struct {
wantlist map[u.Key]struct{} wantlist map[u.Key]struct{}
blocks []blocks.Block blocks map[u.Key]blocks.Block
} }
func New() BitSwapMessage { func New() BitSwapMessage {
return &impl{ return &impl{
wantlist: make(map[u.Key]struct{}), wantlist: make(map[u.Key]struct{}),
blocks: make(map[u.Key]blocks.Block),
} }
} }
...@@ -61,7 +61,11 @@ func (m *impl) Wantlist() []u.Key { ...@@ -61,7 +61,11 @@ func (m *impl) Wantlist() []u.Key {
// TODO(brian): convert these into blocks // TODO(brian): convert these into blocks
func (m *impl) Blocks() []blocks.Block { func (m *impl) Blocks() []blocks.Block {
return m.blocks bs := make([]blocks.Block, 0)
for _, block := range m.blocks {
bs = append(bs, block)
}
return bs
} }
func (m *impl) AddWanted(k u.Key) { func (m *impl) AddWanted(k u.Key) {
...@@ -69,7 +73,7 @@ func (m *impl) AddWanted(k u.Key) { ...@@ -69,7 +73,7 @@ func (m *impl) AddWanted(k u.Key) {
} }
func (m *impl) AppendBlock(b blocks.Block) { func (m *impl) AppendBlock(b blocks.Block) {
m.blocks = append(m.blocks, b) m.blocks[b.Key()] = b
} }
func FromNet(nmsg netmsg.NetMessage) (BitSwapMessage, error) { func FromNet(nmsg netmsg.NetMessage) (BitSwapMessage, error) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment