github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/convert/pprof/streaming/profile_finder.go (about)

     1  package streaming
     2  
     3  import (
     4  	"sort"
     5  )
     6  
     7  func newFinder(functions []function, locations []location) finder {
     8  	res := finder{functions: functions, locations: locations}
     9  	if !locationSlice(locations) {
    10  		res.locationsMap = locationMap(locations)
    11  	}
    12  	if !functionSlice(functions) {
    13  		res.functionsMap = functionMap(functions)
    14  	}
    15  	return res
    16  }
    17  
    18  type finder struct {
    19  	functions    []function
    20  	locations    []location
    21  	functionsMap map[uint64]*function
    22  	locationsMap map[uint64]*location
    23  }
    24  
    25  func (f *finder) FindLocation(id uint64) (*location, bool) {
    26  	if f.locationsMap == nil {
    27  		idx := id - 1
    28  		return &f.locations[idx], true
    29  	}
    30  	l, ok := f.locationsMap[id]
    31  	return l, ok
    32  }
    33  
    34  func (f *finder) FindFunction(id uint64) (*function, bool) {
    35  	if f.functionsMap == nil {
    36  		idx := id - 1
    37  		return &f.functions[idx], true
    38  	}
    39  	ff, ok := f.functionsMap[id]
    40  	return ff, ok
    41  }
    42  
    43  func locationSlice(locations []location) (ok bool) {
    44  	// Check if it's already sorted first
    45  	max := uint64(0)
    46  	sorted := true
    47  	for i, l := range locations {
    48  		if l.id != uint64(i+1) {
    49  			sorted = false
    50  			if l.id > max {
    51  				max = l.id
    52  			}
    53  		}
    54  	}
    55  	if max > uint64(len(locations)) {
    56  		// IDs are not consecutive numbers starting at 1, a slice is not good enough
    57  		return false
    58  	}
    59  	if !sorted {
    60  		sort.Slice(locations, func(i, j int) bool {
    61  			return locations[i].id < locations[j].id
    62  		})
    63  	}
    64  	return true
    65  }
    66  
    67  func locationMap(locations []location) map[uint64]*location {
    68  	m := make(map[uint64]*location, len(locations))
    69  	for i := range locations {
    70  		m[locations[i].id] = &locations[i]
    71  	}
    72  	return m
    73  }
    74  
    75  func functionSlice(functions []function) (ok bool) {
    76  	// Check if it's already sorted first
    77  	max := uint64(0)
    78  	sorted := true
    79  	for i, f := range functions {
    80  		if f.id != uint64(i+1) {
    81  			sorted = false
    82  			if f.id > max {
    83  				max = f.id
    84  			}
    85  		}
    86  	}
    87  	if max > uint64(len(functions)) {
    88  		// IDs are not consecutive numbers starting at one, this won't work
    89  		return false
    90  	}
    91  	if !sorted {
    92  		sort.Slice(functions, func(i, j int) bool {
    93  			return functions[i].id < functions[j].id
    94  		})
    95  	}
    96  	return true
    97  }
    98  
    99  func functionMap(functions []function) map[uint64]*function {
   100  	m := make(map[uint64]*function, len(functions))
   101  	for i := range functions {
   102  		m[functions[i].id] = &functions[i]
   103  	}
   104  	return m
   105  }