github.com/anonymouse64/snapd@v0.0.0-20210824153203-04c4c42d842d/strutil/matchcounter.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2018 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package strutil
    21  
    22  import (
    23  	"bytes"
    24  	"regexp"
    25  )
    26  
    27  // A MatchCounter is a discarding io.Writer that retains up to N
    28  // matches to its Regexp before just counting matches.
    29  //
    30  // It does not work with regexps that cross newlines; in fact it will
    31  // probably not work if the data written isn't line-orineted.
    32  //
    33  // If Regexp is not set (or nil), it matches whole non-empty lines.
    34  type MatchCounter struct {
    35  	// Regexp to use to find matches in the stream
    36  	Regexp *regexp.Regexp
    37  	// Maximum number of matches to keep; if < 0, keep all matches
    38  	N int
    39  	// LastN when true indicates that only N last matches shall be kept.
    40  	LastN bool
    41  
    42  	count   int
    43  	matches []string
    44  	partial bytes.Buffer
    45  }
    46  
    47  func (w *MatchCounter) Write(p []byte) (int, error) {
    48  	n := len(p)
    49  	if w.partial.Len() > 0 {
    50  		idx := bytes.IndexByte(p, '\n')
    51  		if idx < 0 {
    52  			// no newline yet, carry on accumulating
    53  			w.partial.Write(p)
    54  			return n, nil
    55  		}
    56  		idx++
    57  		w.partial.Write(p[:idx])
    58  		w.check(w.partial.Bytes())
    59  		p = p[idx:]
    60  	}
    61  	w.partial.Reset()
    62  	idx := bytes.LastIndexByte(p, '\n')
    63  	if idx < 0 {
    64  		w.partial.Write(p)
    65  		return n, nil
    66  	}
    67  	idx++
    68  	w.partial.Write(p[idx:])
    69  	w.check(p[:idx])
    70  	return n, nil
    71  }
    72  
    73  func (w *MatchCounter) check(p []byte) {
    74  	addMatch := func(m string) {
    75  		switch {
    76  		case w.N == 0:
    77  			// keep none
    78  			return
    79  		case w.N < 0:
    80  			// keep all
    81  			fallthrough
    82  		case !w.LastN && len(w.matches) < w.N:
    83  			w.matches = append(w.matches, m)
    84  		case w.LastN:
    85  			// keep only last N matches
    86  			keep := w.matches
    87  			if len(w.matches)+1 > w.N {
    88  				keep = w.matches[1:]
    89  			}
    90  			w.matches = append(keep, m)
    91  		}
    92  	}
    93  	if w.Regexp == nil {
    94  		for {
    95  			idx := bytes.IndexByte(p, '\n')
    96  			if idx < 0 {
    97  				return
    98  			}
    99  			if idx == 0 {
   100  				// empty line
   101  				p = p[1:]
   102  				continue
   103  			}
   104  			addMatch(string(p[:idx]))
   105  			w.count++
   106  			p = p[idx+1:]
   107  		}
   108  	}
   109  	matches := w.Regexp.FindAll(p, -1)
   110  	for _, match := range matches {
   111  		addMatch(string(match))
   112  	}
   113  	w.count += len(matches)
   114  }
   115  
   116  // Matches returns the first few matches, and the total number of matches seen.
   117  func (w *MatchCounter) Matches() ([]string, int) {
   118  	return w.matches, w.count
   119  }