diff --git a/stream/mts/psi/psi.go b/stream/mts/psi/psi.go
index 11e36d38..f1033be7 100644
--- a/stream/mts/psi/psi.go
+++ b/stream/mts/psi/psi.go
@@ -239,6 +239,11 @@ func asByte(b bool) byte {
 	return 0x00
 }
 
+// AddDescriptor adds or updates a descriptor in a PSI given a descriptor tag
+// and data. If the psi is not a pmt, then an error is returned. If a descriptor
+// with the given tag is not found in the psi, room is made and a descriptor with
+// given tag and data is created. If a descriptor with the tag is found, the
+// descriptor is resized as required and the new data is copied in.
 func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 	if psi.TableID(*p) != pmtID {
 		return errors.New("trying to add descriptor, but not pmt")
@@ -250,27 +255,26 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 		return nil
 	}
 
-	// Old lengths
 	oldDescLen := desc.len()
 	oldDataLen := int(desc[1])
-
-	// New lengths
 	newDataLen := len(data)
 	newDescLen := 2 + newDataLen
-
 	delta := newDescLen - oldDescLen
 
-	if oldDataLen > newDataLen {
+	// If the old data length is more than the new data length, we need shift data
+	// after descriptor up, and then trim the psi. If the oldDataLen is less than
+	// new data then we need reseize psi and shift data down. If data is same size
+	// just copy new data in.
+	switch {
+	case oldDataLen > newDataLen:
 		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
 		*p = (*p)[:len(*p)+delta]
-
-	} else if oldDataLen < newDataLen {
+	case oldDataLen < newDataLen:
 		tmp := make([]byte, len(*p)+delta)
 		copy(tmp, *p)
 		*p = tmp
 		copy((*p)[i+newDescLen:], (*p)[i+oldDescLen:])
-
-	} else {
+	default:
 		copy((*p)[i+2:], data)
 	}
 
@@ -278,20 +282,27 @@ func (p *PSIBytes) AddDescriptor(tag int, data []byte) error {
 	p.setProgInfoLen(newProgInfoLen)
 	newSectionLen := int(psi.SectionLength(*p)) + delta
 	p.setSectionLen(newSectionLen)
-
 	updateCrc((*p)[1:])
-
 	return nil
 }
 
+// len returns the length of a descriptor in bytes.
 func (d *Descriptor) len() int {
 	return int(2 + (*d)[1])
 }
 
+// ProgramInfoLen returns the program info length of a PSI.
+//
+// TODO: check if pmt - if not return 0 ? or -1 ?
 func (p *PSIBytes) ProgramInfoLen() int {
 	return int((((*p)[ProgramInfoLenIdx1] & ProgramInfoLenMask1) << 8) | (*p)[ProgramInfoLenIdx2])
 }
 
+// HasDescriptor checks if a descriptor of the given tag exists in a PSI. If the descriptor
+// of the given tag exists, an index of this descriptor, as well as the Descriptor is returned.
+// If the descriptor of the given tag cannot be found, -1 and a nil slice is returned.
+//
+// TODO: check if pmt, return error if not ?
 func (p *PSIBytes) HasDescriptor(tag int) (int, Descriptor) {
 	descs := p.descriptors()
 	if descs == nil {
@@ -305,20 +316,21 @@ func (p *PSIBytes) HasDescriptor(tag int) (int, Descriptor) {
 	return -1, nil
 }
 
+// descriptors returns the descriptors in a psi if they exist, otherwise
+// a nil slice is returned.
 func (p *PSIBytes) descriptors() []byte {
 	return (*p)[DescriptorsIdx : DescriptorsIdx+p.ProgramInfoLen()]
 }
 
+// createDescriptor creates a descriptor in a psi given a tag and data.
 func (p *PSIBytes) createDescriptor(tag int, data []byte) {
 	curProgLen := p.ProgramInfoLen()
 	oldSyntaxSectionLen := SyntaxSecLenFrom(*p)
 	dataLen := len(data)
-
-	// Calculate the new descriptors index and length.
 	newDescIdx := DescriptorsIdx + curProgLen
 	newDescLen := dataLen + 2
 
-	// Copy data down from newDescIdx to create room for the new descriptor.
+	// Increase size of psi and copy data down to make room for new descriptor.
 	tmp := make([]byte, len(*p)+newDescLen)
 	copy(tmp, *p)
 	*p = tmp
@@ -329,25 +341,23 @@ func (p *PSIBytes) createDescriptor(tag int, data []byte) {
 	(*p)[newDescIdx+1] = byte(dataLen)
 	copy((*p)[newDescIdx+2:newDescIdx+2+dataLen], data)
 
-	// Update the program info length to account for the new descriptor.
+	// Set length fields and update the psi crc.
 	addedLen := dataLen + 2
 	newProgInfoLen := curProgLen + addedLen
 	p.setProgInfoLen(newProgInfoLen)
-
-	// set section length
 	newSyntaxSectionLen := int(oldSyntaxSectionLen) + addedLen
 	p.setSectionLen(newSyntaxSectionLen)
-
 	updateCrc((*p)[1:])
 }
 
+// setProgInfoLen sets the program information length in a psi with a pmt.
 func (p *PSIBytes) setProgInfoLen(l int) {
-	// TODO: check if pmt first
 	(*p)[ProgramInfoLenIdx1] &= 0xff ^ ProgramInfoLenMask1
 	(*p)[ProgramInfoLenIdx1] |= byte(l>>8) & ProgramInfoLenMask1
 	(*p)[ProgramInfoLenIdx2] = byte(l)
 }
 
+// setSectionLen sets section length in a psi.
 func (p *PSIBytes) setSectionLen(l int) {
 	(*p)[SyntaxSecLenIdx1] &= 0xff ^ SyntaxSecLenMask1
 	(*p)[SyntaxSecLenIdx1] |= byte(l>>8) & SyntaxSecLenMask1