segment: Use BindTo and Message exclusively

In the future, this will make it possible external implementations of
Segment to be written.

It also fixes the final data race in the refactor.

arena.go

@@ -106,7 +106,7 @@ func (ssa *SingleSegmentArena) Allocate(sz Size, msg *Message, seg *Segment) (*S
 	if len(data)%int(wordSize) != 0 {
 		return nil, 0, errors.New("segment size is not a multiple of word size")
 	}
-	ssa.seg.msg = msg
+	ssa.seg.BindTo(msg)
 	if hasCapacity(data, sz) {
 		addr := address(len(ssa.seg.data))
 		ssa.seg.data = ssa.seg.data[:len(ssa.seg.data)+int(sz)]
@@ -145,7 +145,7 @@ func (ssa *SingleSegmentArena) Release() {
 		zeroSlice(ssa.seg.data)
 		ssa.bp.Put(ssa.seg.data)
 	}
-	ssa.seg.msg = nil
+	ssa.seg.BindTo(nil)
 	ssa.seg.data = nil
 	if ssa.fromPool {
 		ssa.fromPool = false // Prevent double return
@@ -196,7 +196,7 @@ func (msa *MultiSegmentArena) Release() {
 			msa.bp.Put(msa.segs[i].data)
 		}
 		msa.segs[i].data = nil
-		msa.segs[i].msg = nil
+		msa.segs[i].BindTo(nil)
 	}
 
 	if msa.segs != nil {
@@ -293,14 +293,14 @@ func (msa *MultiSegmentArena) Allocate(sz Size, msg *Message, seg *Segment) (*Se
 		}
 
 		// Double check this segment is for this message.
-		if seg.msg != nil && seg.msg != msg {
+		if seg.Message() != nil && seg.Message() != msg {
 			return nil, 0, errors.New("attempt to allocate in segment for different message")
 		}
 
 		addr := address(len(seg.data))
 		newLen := int(addr) + int(sz)
 		seg.data = seg.data[:newLen]
-		seg.msg = msg
+		seg.BindTo(msg)
 		return seg, addr, nil
 	}
 
@@ -312,7 +312,7 @@ func (msa *MultiSegmentArena) Allocate(sz Size, msg *Message, seg *Segment) (*Se
 			addr := address(len(msa.segs[i].data))
 			newLen := int(addr) + int(sz)
 			msa.segs[i].data = msa.segs[i].data[:newLen]
-			msa.segs[i].msg = msg
+			msa.segs[i].BindTo(msg)
 			return &msa.segs[i], addr, nil
 		}
 
@@ -354,9 +354,10 @@ func (msa *MultiSegmentArena) Allocate(sz Size, msg *Message, seg *Segment) (*Se
 	msa.segs = append(msa.segs, Segment{
 		data: buf,
 		id:   id,
-		msg:  msg,
 	})
-	return &msa.segs[int(id)], 0, nil
+	res := &msa.segs[int(id)]
+	res.BindTo(msg)
+	return res, 0, nil
 }
 
 func (msa *MultiSegmentArena) String() string {

capability.go

@@ -59,7 +59,7 @@ func (i Interface) Message() *Message {
 	if i.seg == nil {
 		return nil
 	}
-	return i.seg.msg
+	return i.seg.Message()
 }
 
 // IsValid returns whether the interface is valid.

list.go

@@ -100,7 +100,7 @@ func (p List) Message() *Message {
 	if p.seg == nil {
 		return nil
 	}
-	return p.seg.msg
+	return p.seg.Message()
 }
 
 // IsValid returns whether the list is valid.

message.go

@@ -256,10 +256,12 @@ func (m *Message) Segment(id SegmentID) (*Segment, error) {
 	if seg == nil {
 		return nil, errors.New("segment " + str.Utod(id) + " out of bounds in arena")
 	}
-	if seg.msg != nil && seg.msg != m {
+	segMsg := seg.Message()
+	if segMsg == nil {
+		seg.BindTo(m)
+	} else if segMsg != m {
 		return nil, errors.New("segment " + str.Utod(id) + ": not of the same message")
 	}
-	seg.msg = m
 	return seg, nil
 }
 
@@ -352,17 +354,18 @@ func alloc(s *Segment, sz Size) (*Segment, address, error) {
 	}
 	sz = sz.padToWord()
 
-	if s.msg == nil {
+	msg := s.Message()
+	if msg == nil {
 		return nil, 0, errors.New("segment does not have a message assotiated with it")
 	}
-	if s.msg.Arena == nil {
+	if msg.Arena == nil {
 		return nil, 0, errors.New("message does not have an arena")
 	}
 
 	// TODO: From this point on, this could be changed to be a requirement
 	// for Arena implementations instead of relying on alloc() to do it.
 
-	s, addr, err := s.msg.Arena.Allocate(sz, s.msg, s)
+	s, addr, err := msg.Arena.Allocate(sz, msg, s)
 	if err != nil {
 		return s, addr, err
 	}

pointer.go

@@ -187,7 +187,7 @@ func (p Ptr) Message() *Message {
 	if p.seg == nil {
 		return nil
 	}
-	return p.seg.msg
+	return p.seg.Message()
 }
 
 // Default returns p if it is valid, otherwise it unmarshals def.

segment.go

@@ -17,7 +17,8 @@ type SegmentID uint32
 type Segment struct {
 	// msg associated with this segment. A Message instance m maintains the
 	// invariant that that all m.segs[].msg == m.
-	msg  *Message
+	msg *Message
+
 	id   SegmentID
 	data []byte
 }
@@ -27,6 +28,12 @@ func (s *Segment) Message() *Message {
 	return s.msg
 }
 
+// BindTo binds the segment to a given message. This is usually only called by
+// Arena implementations and does not perform any kind of safety check.
+func (s *Segment) BindTo(m *Message) {
+	s.msg = m
+}
+
 // ID returns the segment's ID.
 func (s *Segment) ID() SegmentID {
 	return s.id
@@ -107,15 +114,15 @@ func (s *Segment) root() PointerList {
 		seg:        s,
 		length:     1,
 		size:       sz,
-		depthLimit: s.msg.depthLimit(),
+		depthLimit: s.Message().depthLimit(),
 	}
 }
 
 func (s *Segment) lookupSegment(id SegmentID) (*Segment, error) {
 	if s.id == id {
 		return s, nil
 	}
-	return s.msg.Segment(id)
+	return s.Message().Segment(id)
 }
 
 func (s *Segment) readPtr(paddr address, depthLimit uint) (ptr Ptr, err error) {
@@ -135,7 +142,7 @@ func (s *Segment) readPtr(paddr address, depthLimit uint) (ptr Ptr, err error) {
 		if err != nil {
 			return Ptr{}, exc.WrapError("read pointer", err)
 		}
-		if !s.msg.canRead(sp.readSize()) {
+		if !s.Message().canRead(sp.readSize()) {
 			return Ptr{}, errors.New("read pointer: read traversal limit reached")
 		}
 		sp.depthLimit = depthLimit - 1
@@ -145,7 +152,7 @@ func (s *Segment) readPtr(paddr address, depthLimit uint) (ptr Ptr, err error) {
 		if err != nil {
 			return Ptr{}, exc.WrapError("read pointer", err)
 		}
-		if !s.msg.canRead(lp.readSize()) {
+		if !s.Message().canRead(lp.readSize()) {
 			return Ptr{}, errors.New("read pointer: read traversal limit reached")
 		}
 		lp.depthLimit = depthLimit - 1
@@ -377,8 +384,8 @@ func (s *Segment) writePtr(off address, src Ptr, forceCopy bool) error {
 		srcRaw = l.raw()
 	case interfacePtrType:
 		i := src.Interface()
-		if src.seg.msg != s.msg {
-			c := s.msg.CapTable().Add(i.Client().AddRef())
+		if src.seg.Message() != s.Message() {
+			c := s.Message().CapTable().Add(i.Client().AddRef())
 			i = NewInterface(s, c)
 		}
 		s.writeRawPointer(off, i.value(off))

struct.go

@@ -46,7 +46,7 @@ func NewRootStruct(s *Segment, sz ObjectSize) (Struct, error) {
 	if err != nil {
 		return st, err
 	}
-	if err := s.msg.SetRoot(st.ToPtr()); err != nil {
+	if err := s.Message().SetRoot(st.ToPtr()); err != nil {
 		return st, err
 	}
 	return st, nil
@@ -75,7 +75,7 @@ func (p Struct) Message() *Message {
 	if p.seg == nil {
 		return nil
 	}
-	return p.seg.msg
+	return p.seg.Message()
 }
 
 // IsValid returns whether the struct is valid.