cleanups

asticode · Apr 14, 2021 · d1f8a96 · d1f8a96
1 parent 64c4e19
commit d1f8a96
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/data_psi.go b/data_psi.go
@@ -119,8 +119,7 @@ func parsePSIData(i *astikit.BytesIterator) (d *PSIData, err error) {
 		if s, stop, err = parsePSISection(i); err != nil {
 			err = fmt.Errorf("astits: parsing PSI table failed: %w", err)
 			return
-		}
-		if stop {
+		} else if stop {
 			break
 		}
 		d.Sections = append(d.Sections, s)
@@ -135,14 +134,10 @@ func parsePSISection(i *astikit.BytesIterator) (s *PSISection, stop bool, err er
 
 	// Parse header
 	var offsetStart, offsetSectionsEnd, offsetEnd int
-	if s.Header, offsetStart, _, offsetSectionsEnd, offsetEnd, err = parsePSISectionHeader(i); err != nil {
+	if s.Header, offsetStart, _, offsetSectionsEnd, offsetEnd, stop, err = parsePSISectionHeader(i); err != nil {
 		err = fmt.Errorf("astits: parsing PSI section header failed: %w", err)
 		return
-	}
-
-	// Check whether we need to stop the parsing
-	if shouldStopPSIParsing(s.Header.TableID) {
-		stop = true
+	} else if stop {
 		return
 	}
 
@@ -206,7 +201,7 @@ func shouldStopPSIParsing(tableID PSITableID) bool {
 }
 
 // parsePSISectionHeader parses a PSI section header
-func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd int, err error) {
+func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd int, stop bool, err error) {
 	// Init
 	h = &PSISectionHeader{}
 	offsetStart = i.Offset()
@@ -225,7 +220,8 @@ func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offse
 	h.TableType = h.TableID.Type()
 
 	// Check whether we need to stop the parsing
-	if shouldStopPSIParsing(h.TableID) {
+	if h.TableID == PSITableIDNull {
+		stop = true
 		return
 	}
 

diff --git a/data_psi_test.go b/data_psi_test.go
@@ -212,20 +212,21 @@ func TestParsePSISectionHeader(t *testing.T) {
 	w.Write(uint8(254)) // Table ID
 	w.Write("1")        // Syntax section indicator
 	w.Write("0000000")  // Finish the byte
-	d, _, _, _, _, err := parsePSISectionHeader(astikit.NewBytesIterator(buf.Bytes()))
+	d, _, _, _, _, _, err := parsePSISectionHeader(astikit.NewBytesIterator(buf.Bytes()))
 	assert.Equal(t, d, &PSISectionHeader{
 		TableID:   254,
 		TableType: PSITableTypeUnknown,
 	})
 	assert.NoError(t, err)
 
 	// Valid table type
-	d, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd, err := parsePSISectionHeader(astikit.NewBytesIterator(psiSectionHeaderBytes()))
+	d, offsetStart, offsetSectionsStart, offsetSectionsEnd, offsetEnd, stop, err := parsePSISectionHeader(astikit.NewBytesIterator(psiSectionHeaderBytes()))
 	assert.Equal(t, d, psiSectionHeader)
 	assert.Equal(t, 0, offsetStart)
 	assert.Equal(t, 3, offsetSectionsStart)
 	assert.Equal(t, 2729, offsetSectionsEnd)
 	assert.Equal(t, 2733, offsetEnd)
+	assert.Equal(t, false, stop)
 	assert.NoError(t, err)
 }