1// Copyright 2017 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// Package remap handles tracking the locations of Go tokens in a source text
6// across a rewrite by the Go formatter.
7package remap
8
9import (
10	"fmt"
11	"go/scanner"
12	"go/token"
13)
14
15// A Location represents a span of byte offsets in the source text.
16type Location struct {
17	Pos, End int // End is exclusive
18}
19
20// A Map represents a mapping between token locations in an input source text
21// and locations in the correspnding output text.
22type Map map[Location]Location
23
24// Find reports whether the specified span is recorded by m, and if so returns
25// the new location it was mapped to. If the input span was not found, the
26// returned location is the same as the input.
27func (m Map) Find(pos, end int) (Location, bool) {
28	key := Location{
29		Pos: pos,
30		End: end,
31	}
32	if loc, ok := m[key]; ok {
33		return loc, true
34	}
35	return key, false
36}
37
38func (m Map) add(opos, oend, npos, nend int) {
39	m[Location{Pos: opos, End: oend}] = Location{Pos: npos, End: nend}
40}
41
42// Compute constructs a location mapping from input to output.  An error is
43// reported if any of the tokens of output cannot be mapped.
44func Compute(input, output []byte) (Map, error) {
45	itok := tokenize(input)
46	otok := tokenize(output)
47	if len(itok) != len(otok) {
48		return nil, fmt.Errorf("wrong number of tokens, %d ≠ %d", len(itok), len(otok))
49	}
50	m := make(Map)
51	for i, ti := range itok {
52		to := otok[i]
53		if ti.Token != to.Token {
54			return nil, fmt.Errorf("token %d type mismatch: %s ≠ %s", i+1, ti, to)
55		}
56		m.add(ti.pos, ti.end, to.pos, to.end)
57	}
58	return m, nil
59}
60
61// tokinfo records the span and type of a source token.
62type tokinfo struct {
63	pos, end int
64	token.Token
65}
66
67func tokenize(src []byte) []tokinfo {
68	fs := token.NewFileSet()
69	var s scanner.Scanner
70	s.Init(fs.AddFile("src", fs.Base(), len(src)), src, nil, scanner.ScanComments)
71	var info []tokinfo
72	for {
73		pos, next, lit := s.Scan()
74		switch next {
75		case token.SEMICOLON:
76			continue
77		}
78		info = append(info, tokinfo{
79			pos:   int(pos - 1),
80			end:   int(pos + token.Pos(len(lit)) - 1),
81			Token: next,
82		})
83		if next == token.EOF {
84			break
85		}
86	}
87	return info
88}
89