mask.go

   1 package lightning
   2
   3 import (
   4         "sort"
   5 )
   6
   7 type interval struct {
   8         start int
   9         end   int
  10 }
  11
  12 type intervalTreeNode struct {
  13         interval interval
  14         maxend   int
  15 }
  16
  17 type intervalTree []intervalTreeNode
  18
  19 type mask struct {
  20         intervals map[string][]interval
  21         itrees    map[string]intervalTree
  22         frozen    bool
  23 }
  24
  25 func (m *mask) Add(seqname string, start, end int) {
  26         if m.intervals == nil {
  27                 m.intervals = map[string][]interval{}
  28         }
  29         m.intervals[seqname] = append(m.intervals[seqname], interval{start, end})
  30 }
  31
  32 func (m *mask) Freeze() {
  33         m.itrees = map[string]intervalTree{}
  34         for seqname, intervals := range m.intervals {
  35                 m.itrees[seqname] = m.freeze(intervals)
  36         }
  37         m.frozen = true
  38 }
  39
  40 func (m *mask) Check(seqname string, start, end int) bool {
  41         if !m.frozen {
  42                 panic("bug: (*mask)Check() called before Freeze()")
  43         }
  44         return m.itrees[seqname].check(0, interval{start, end})
  45 }
  46
  47 func (m *mask) freeze(in []interval) intervalTree {
  48         if len(in) == 0 {
  49                 return nil
  50         }
  51         sort.Slice(in, func(i, j int) bool {
  52                 return in[i].start < in[j].start
  53         })
  54         itreesize := 1
  55         for itreesize < len(in) {
  56                 itreesize = itreesize * 2
  57         }
  58         itree := make(intervalTree, itreesize)
  59         itree.importSlice(0, in)
  60         for i := len(in); i < itreesize; i++ {
  61                 itree[i].maxend = -1
  62         }
  63         return itree
  64 }
  65
  66 func (itree intervalTree) check(root int, q interval) bool {
  67         return root < len(itree) &&
  68                 itree[root].maxend >= q.start &&
  69                 ((itree[root].interval.start <= q.end && itree[root].interval.end >= q.start) ||
  70                         itree.check(root*2+1, q) ||
  71                         itree.check(root*2+2, q))
  72 }
  73
  74 func (itree intervalTree) importSlice(root int, in []interval) int {
  75         mid := len(in) / 2
  76         node := intervalTreeNode{interval: in[mid], maxend: in[mid].end}
  77         if mid > 0 {
  78                 end := itree.importSlice(root*2+1, in[0:mid])
  79                 if end > node.maxend {
  80                         node.maxend = end
  81                 }
  82         }
  83         if mid+1 < len(in) {
  84                 end := itree.importSlice(root*2+2, in[mid+1:])
  85                 if end > node.maxend {
  86                         node.maxend = end
  87                 }
  88         }
  89         itree[root] = node
  90         return node.maxend
  91 }