rewrite.mx raw

   1  package rewrite
   2  
   3  import "bytes"
   4  
   5  func StripDuplicatePackageClauses(src []byte) []byte {
   6  	found := false
   7  	var out []byte
   8  	i := 0
   9  	for i < len(src) {
  10  		nlIdx := bytes.IndexByte(src[i:], '\n')
  11  		var line []byte
  12  		var lineEnd int32
  13  		if nlIdx < 0 {
  14  			line = src[i:]
  15  			lineEnd = len(src)
  16  		} else {
  17  			line = src[i : i+nlIdx]
  18  			lineEnd = i + nlIdx + 1
  19  		}
  20  		trimmed := bytes.TrimSpace(line)
  21  		if bytes.HasPrefix(trimmed, "package ") {
  22  			if found {
  23  				if out == nil {
  24  					out = []byte{:0:len(src)}
  25  					out = append(out, src[:i]...)
  26  				}
  27  				for k := 0; k < len(line); k++ {
  28  					out = append(out, ' ')
  29  				}
  30  				if nlIdx >= 0 {
  31  					out = append(out, '\n')
  32  				}
  33  				i = lineEnd
  34  				continue
  35  			}
  36  			found = true
  37  		}
  38  		if out != nil {
  39  			out = append(out, src[i:lineEnd]...)
  40  		}
  41  		i = lineEnd
  42  	}
  43  	if out == nil {
  44  		return src
  45  	}
  46  	return out
  47  }
  48  
  49  func RewriteSliceMakeLiterals(src []byte) []byte {
  50  	var out []byte
  51  	i := 0
  52  	for i < len(src) {
  53  		start := bytes.Index(src[i:], []byte("{:"))
  54  		if start < 0 {
  55  			out = append(out, src[i:]...)
  56  			break
  57  		}
  58  		start = start + i
  59  		lbrack := findSliceTypeStart(src, start)
  60  		if lbrack < 0 {
  61  			out = append(out, src[i:start+2]...)
  62  			i = start + 2
  63  			continue
  64  		}
  65  		close := findMatchingBrace(src, start)
  66  		if close < 0 {
  67  			out = append(out, src[i:start+2]...)
  68  			i = start + 2
  69  			continue
  70  		}
  71  		inner := src[start+2 : close]
  72  		colonIdx := bytes.IndexByte(inner, ':')
  73  		typeText := src[lbrack:start]
  74  		out = append(out, src[i:lbrack]...)
  75  		if colonIdx < 0 {
  76  			out = append(out, "make("...)
  77  			out = append(out, typeText...)
  78  			out = append(out, ", "...)
  79  			out = append(out, bytes.TrimSpace(inner)...)
  80  			out = append(out, ')')
  81  		} else {
  82  			lenExpr := bytes.TrimSpace(inner[:colonIdx])
  83  			capExpr := bytes.TrimSpace(inner[colonIdx+1:])
  84  			out = append(out, "make("...)
  85  			out = append(out, typeText...)
  86  			out = append(out, ", "...)
  87  			out = append(out, lenExpr...)
  88  			out = append(out, ", "...)
  89  			out = append(out, capExpr...)
  90  			out = append(out, ')')
  91  		}
  92  		i = close + 1
  93  	}
  94  	if out == nil {
  95  		return src
  96  	}
  97  	return out
  98  }
  99  
 100  func findSliceTypeStart(src []byte, braceIdx int32) int32 {
 101  	j := braceIdx - 1
 102  	for j >= 0 && (src[j] == ' ' || src[j] == '\t' || src[j] == '\n') {
 103  		j--
 104  	}
 105  	if j < 0 {
 106  		return -1
 107  	}
 108  	depth := 0
 109  	parenDepth := 0
 110  	candidate := -1
 111  	for j >= 0 {
 112  		ch := src[j]
 113  		if ch == ')' {
 114  			parenDepth++
 115  		} else if ch == '(' {
 116  			if parenDepth == 0 {
 117  				if candidate >= 0 {
 118  					return candidate
 119  				}
 120  				return -1
 121  			}
 122  			parenDepth--
 123  		} else if parenDepth > 0 {
 124  			j--
 125  			continue
 126  		}
 127  		if ch == ']' {
 128  			depth++
 129  		} else if ch == '[' {
 130  			depth--
 131  			if depth == 0 {
 132  				candidate = j
 133  			}
 134  		} else if ch == '{' || ch == ';' {
 135  			if candidate >= 0 {
 136  				return candidate
 137  			}
 138  			return -1
 139  		} else if depth == 0 && parenDepth == 0 {
 140  			if candidate >= 0 {
 141  				return candidate
 142  			}
 143  			if ch == ' ' || ch == '\t' || ch == '\n' || ch == '*' || ch == '(' || ch == ')' {
 144  				j--
 145  				continue
 146  			}
 147  			if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_' || ch == '.' {
 148  				j--
 149  				continue
 150  			}
 151  			return -1
 152  		}
 153  		j--
 154  	}
 155  	if candidate >= 0 {
 156  		return candidate
 157  	}
 158  	return -1
 159  }
 160  
 161  func findMatchingBrace(src []byte, openIdx int32) int32 {
 162  	depth := 1
 163  	for i := openIdx + 1; i < len(src); i++ {
 164  		if src[i] == '{' {
 165  			depth++
 166  		} else if src[i] == '}' {
 167  			depth--
 168  			if depth == 0 {
 169  				return i
 170  			}
 171  		}
 172  	}
 173  	return -1
 174  }
 175  
 176  func RewriteChanMakeLiterals(src []byte) []byte {
 177  	chanKw := []byte("chan ")
 178  	var out []byte
 179  	i := 0
 180  	for i < len(src) {
 181  		idx := bytes.Index(src[i:], chanKw)
 182  		if idx < 0 {
 183  			if out != nil {
 184  				out = append(out, src[i:]...)
 185  			}
 186  			break
 187  		}
 188  		idx = idx + i
 189  		j := idx + 5
 190  		for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
 191  			j++
 192  		}
 193  		if j >= len(src) {
 194  			if out != nil {
 195  				out = append(out, src[i:]...)
 196  			}
 197  			break
 198  		}
 199  		if src[j] == '<' && j+1 < len(src) && src[j+1] == '-' {
 200  			if out != nil {
 201  				out = append(out, src[i:j+2]...)
 202  			}
 203  			i = j + 2
 204  			continue
 205  		}
 206  		for j < len(src) && (src[j] != '{' && src[j] != '\n' && src[j] != ';' && src[j] != '(' && src[j] != ')') {
 207  			if src[j] == ' ' || src[j] == '\t' {
 208  				break
 209  			}
 210  			j++
 211  		}
 212  		if j >= len(src) || src[j] != '{' {
 213  			if out == nil {
 214  				i = idx + 4
 215  			} else {
 216  				out = append(out, src[i:idx+4]...)
 217  				i = idx + 4
 218  			}
 219  			continue
 220  		}
 221  		braceOpen := j
 222  		endsStruct := braceOpen >= 6 && string(src[braceOpen-6:braceOpen]) == "struct"
 223  		if !endsStruct && braceOpen >= 7 {
 224  			k := braceOpen - 1
 225  			for k > idx && (src[k] == ' ' || src[k] == '\t') {
 226  				k--
 227  			}
 228  			if k >= 5 && string(src[k-5:k+1]) == "struct" {
 229  				endsStruct = true
 230  			}
 231  		}
 232  		if endsStruct {
 233  			structClose := findMatchingBrace(src, braceOpen)
 234  			if structClose < 0 {
 235  				if out == nil {
 236  					i = idx + 4
 237  				} else {
 238  					out = append(out, src[i:idx+4]...)
 239  					i = idx + 4
 240  				}
 241  				continue
 242  			}
 243  			nj := structClose + 1
 244  			if nj >= len(src) || src[nj] != '{' {
 245  				if out == nil {
 246  					i = structClose + 1
 247  				} else {
 248  					out = append(out, src[i:structClose+1]...)
 249  					i = structClose + 1
 250  				}
 251  				continue
 252  			}
 253  			braceOpen = nj
 254  		}
 255  		close := findMatchingBrace(src, braceOpen)
 256  		if close < 0 {
 257  			if out == nil {
 258  				i = idx + 4
 259  			} else {
 260  				out = append(out, src[i:idx+4]...)
 261  				i = idx + 4
 262  			}
 263  			continue
 264  		}
 265  		inner := bytes.TrimSpace(src[braceOpen+1 : close])
 266  		chanType := src[idx : braceOpen]
 267  		for len(chanType) > 0 && (chanType[len(chanType)-1] == ' ' || chanType[len(chanType)-1] == '\t') {
 268  			chanType = chanType[:len(chanType)-1]
 269  		}
 270  		if out == nil {
 271  			out = []byte{:0:len(src)}
 272  			out = append(out, src[:idx]...)
 273  		} else {
 274  			out = append(out, src[i:idx]...)
 275  		}
 276  		if len(inner) == 0 {
 277  			out = append(out, "make("...)
 278  			out = append(out, chanType...)
 279  			out = append(out, ')')
 280  		} else {
 281  			out = append(out, "make("...)
 282  			out = append(out, chanType...)
 283  			out = append(out, ", "...)
 284  			out = append(out, inner...)
 285  			out = append(out, ')')
 286  		}
 287  		i = close + 1
 288  	}
 289  	if out == nil {
 290  		return src
 291  	}
 292  	return out
 293  }
 294