363 lines
9.4 KiB
Go
363 lines
9.4 KiB
Go
|
// CookieJar - A contestant's algorithm toolbox
|
||
|
|
// Copyright 2014 Peter Szilagyi. All rights reserved.
|
||
|
|
//
|
||
|
|
// CookieJar is dual licensed: use of this source code is governed by a BSD
|
||
|
|
// license that can be found in the LICENSE file. Alternatively, the CookieJar
|
||
|
|
// toolbox may be used in accordance with the terms and conditions contained
|
||
|
|
// in a signed written agreement between you and the author(s).
|
||
|
|
|
||
|
|
// Package deps contains utility functions for contest dependency management.
|
||
|
|
package deps
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"go/ast"
|
||
|
|
"go/parser"
|
||
|
|
"go/printer"
|
||
|
|
"go/token"
|
||
|
|
"os"
|
||
|
|
"os/exec"
|
||
|
|
"path/filepath"
|
||
|
|
"regexp"
|
||
|
|
"strings"
|
||
|
|
|
||
|
|
"code.google.com/p/go.tools/imports"
|
||
|
|
"gopkg.in/inconshreveable/log15.v2"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Package description from the go list command
|
||
|
|
type Package struct {
|
||
|
|
Name string
|
||
|
|
Dir string
|
||
|
|
Standard bool
|
||
|
|
Deps []string
|
||
|
|
GoFiles []string
|
||
|
|
}
|
||
|
|
|
||
|
|
// Loads the details of the Go package.
|
||
|
|
func details(name string) (*Package, error) {
|
||
|
|
// Create the command to retrieve the package infos
|
||
|
|
cmd := exec.Command("go", "list", "-e", "-json", name)
|
||
|
|
|
||
|
|
// Retrieve the output, redirect the errors
|
||
|
|
out, err := cmd.StdoutPipe()
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
cmd.Stderr = os.Stderr
|
||
|
|
|
||
|
|
// Start executing and parse the results
|
||
|
|
if err := cmd.Start(); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
defer cmd.Process.Kill()
|
||
|
|
|
||
|
|
info := new(Package)
|
||
|
|
if err := json.NewDecoder(out).Decode(&info); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Clean up and return
|
||
|
|
if err := cmd.Wait(); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return info, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Collects all the imported packages of a file.
|
||
|
|
func dependencies(path string) (map[string][]string, error) {
|
||
|
|
// Retrieve the dependencies of the source file
|
||
|
|
info, err := details(path)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Iterate over every dependency and gather the sources
|
||
|
|
sources := make(map[string][]string)
|
||
|
|
for _, dep := range info.Deps {
|
||
|
|
// Retrieve the dependency details
|
||
|
|
pkgInfo, err := details(dep)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Gather external library sources
|
||
|
|
if !pkgInfo.Standard {
|
||
|
|
for _, src := range pkgInfo.GoFiles {
|
||
|
|
sources[pkgInfo.Name] = append(sources[pkgInfo.Name], filepath.Join(pkgInfo.Dir, src))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return sources, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Iterates over a package contents and collects all declarations to rename.
|
||
|
|
func declarations(paths []string) ([]*ast.Object, error) {
|
||
|
|
results := []*ast.Object{}
|
||
|
|
for _, path := range paths {
|
||
|
|
// Parse the specified source file
|
||
|
|
fileSet := token.NewFileSet()
|
||
|
|
tree, err := parser.ParseFile(fileSet, path, nil, parser.ParseComments)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Collect all top level declarations
|
||
|
|
for _, decl := range tree.Decls {
|
||
|
|
switch decl := decl.(type) {
|
||
|
|
case *ast.FuncDecl:
|
||
|
|
if decl.Recv == nil && decl.Name.Name != "init" {
|
||
|
|
results = append(results, ast.NewObj(ast.Fun, decl.Name.String()))
|
||
|
|
}
|
||
|
|
case *ast.GenDecl:
|
||
|
|
for _, spec := range decl.Specs {
|
||
|
|
switch spec := spec.(type) {
|
||
|
|
case *ast.ValueSpec:
|
||
|
|
for _, name := range spec.Names {
|
||
|
|
results = append(results, ast.NewObj(ast.Var, name.String()))
|
||
|
|
}
|
||
|
|
case *ast.TypeSpec:
|
||
|
|
results = append(results, ast.NewObj(ast.Typ, spec.Name.String()))
|
||
|
|
default:
|
||
|
|
log15.Warn("Unknown specification", "spec", spec)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
log15.Warn("Unknown declaration", "decl", decl)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return results, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Parses a source file and scopes all global declarations.
|
||
|
|
func rewrite(src string, pkg string, decls []*ast.Object) ([]byte, error) {
|
||
|
|
fileSet := token.NewFileSet()
|
||
|
|
tree, err := parser.ParseFile(fileSet, src, nil, parser.ParseComments)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Scope all top level declarations if not main file
|
||
|
|
if pkg != "" {
|
||
|
|
for _, decl := range decls {
|
||
|
|
rename(tree, decl.Name, pkg+"ᴥ"+decl.Name, decl.Kind)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Generate the new source contents
|
||
|
|
out := bytes.NewBuffer(nil)
|
||
|
|
for _, decl := range tree.Decls {
|
||
|
|
if gen, ok := decl.(*ast.GenDecl); !ok || gen.Tok != token.IMPORT {
|
||
|
|
if err := printer.Fprint(out, fileSet, decl); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
fmt.Fprintf(out, "\n\n")
|
||
|
|
}
|
||
|
|
blob := out.Bytes()
|
||
|
|
|
||
|
|
// Scope all externally imported dependencies
|
||
|
|
var fail error
|
||
|
|
ast.Inspect(tree, func(node ast.Node) bool {
|
||
|
|
if imp, ok := node.(*ast.GenDecl); ok && imp.Tok == token.IMPORT {
|
||
|
|
for _, spec := range imp.Specs {
|
||
|
|
if spec, ok := spec.(*ast.ImportSpec); ok {
|
||
|
|
// Figure out the correct name of the import
|
||
|
|
path := strings.Trim(spec.Path.Value, "\"")
|
||
|
|
if info, err := details(path); err != nil {
|
||
|
|
fail = err
|
||
|
|
return false
|
||
|
|
} else if !info.Standard {
|
||
|
|
// Add scope to external import
|
||
|
|
scoper := regexp.MustCompile("([^\\.])" + info.Name + "\\.([^ ])")
|
||
|
|
blob = scoper.ReplaceAll(blob, []byte("${1}"+info.Name+"ᴥ${2}"))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
if fail != nil {
|
||
|
|
return nil, fail
|
||
|
|
}
|
||
|
|
return blob, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Renames a top level declaration to something else.
|
||
|
|
func rename(tree *ast.File, old, new string, kind ast.ObjKind) {
|
||
|
|
// Rename top-level declarations
|
||
|
|
for _, decl := range tree.Decls {
|
||
|
|
switch decl := decl.(type) {
|
||
|
|
case *ast.FuncDecl:
|
||
|
|
// If a top level function matches, rename
|
||
|
|
if decl.Recv == nil && decl.Name.Name == old {
|
||
|
|
decl.Name.Name = new
|
||
|
|
if decl.Name.Obj != nil {
|
||
|
|
decl.Name.Obj.Name = new
|
||
|
|
}
|
||
|
|
}
|
||
|
|
case *ast.GenDecl:
|
||
|
|
// Iterate over all the generic declaration
|
||
|
|
for _, spec := range decl.Specs {
|
||
|
|
switch spec := spec.(type) {
|
||
|
|
case *ast.ValueSpec:
|
||
|
|
// If a top level variable matches, rename
|
||
|
|
for _, name := range spec.Names {
|
||
|
|
if name.Name == old {
|
||
|
|
name.Name = new
|
||
|
|
if name.Obj != nil {
|
||
|
|
name.Obj.Name = new
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
case *ast.TypeSpec:
|
||
|
|
if spec.Name.Name == old {
|
||
|
|
spec.Name.Name = new
|
||
|
|
if spec.Name.Obj != nil {
|
||
|
|
spec.Name.Obj.Name = new
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Walk the AST and rename all internal occurrences
|
||
|
|
stack := []ast.Node{}
|
||
|
|
ast.Inspect(tree, func(node ast.Node) bool {
|
||
|
|
// Keep a traversal stack if need to reference parent
|
||
|
|
if node == nil {
|
||
|
|
stack = stack[:len(stack)-1]
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
stack = append(stack, node)
|
||
|
|
|
||
|
|
// Look for identifiers to rename
|
||
|
|
id, ok := node.(*ast.Ident)
|
||
|
|
if ok && id.Obj == nil && id.Name == old {
|
||
|
|
// Don't rename selected identifiers, member functions, struct keys
|
||
|
|
switch stack[len(stack)-2].(type) {
|
||
|
|
case *ast.SelectorExpr, *ast.FuncDecl:
|
||
|
|
return true
|
||
|
|
case *ast.KeyValueExpr:
|
||
|
|
// If the rename is a variable, allow it
|
||
|
|
if kind != ast.Var && kind != ast.Con {
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
id.Name = new
|
||
|
|
}
|
||
|
|
if ok && id.Obj != nil && id.Name == old && id.Obj.Name == new {
|
||
|
|
// Don't rename struct keys
|
||
|
|
switch stack[len(stack)-2].(type) {
|
||
|
|
case *ast.KeyValueExpr:
|
||
|
|
// If the rename is a variable, allow it
|
||
|
|
if kind != ast.Var && kind != ast.Con {
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
id.Name = new
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// Removes dead code from the source blob.
|
||
|
|
func shake(blob []byte) ([]byte, error) {
|
||
|
|
// Repeat the process until all occurrences have been shaken down
|
||
|
|
cascade := true
|
||
|
|
for cascade {
|
||
|
|
cascade = false
|
||
|
|
|
||
|
|
// Parse the remaining data blob
|
||
|
|
fileSet := token.NewFileSet()
|
||
|
|
tree, err := parser.ParseFile(fileSet, "", blob, 0)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Shake all unused top level declarations
|
||
|
|
retain := []ast.Decl{}
|
||
|
|
for _, decl := range tree.Decls {
|
||
|
|
switch decl := decl.(type) {
|
||
|
|
case *ast.FuncDecl:
|
||
|
|
// If a plain dangling top level function, check usage
|
||
|
|
if name := decl.Name.Name; decl.Recv == nil && strings.Contains(name, "ᴥ") {
|
||
|
|
if bytes.Index(blob, []byte(name)) == bytes.LastIndex(blob, []byte(name)) {
|
||
|
|
cascade = true
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
retain = append(retain, decl)
|
||
|
|
}
|
||
|
|
tree.Decls = retain
|
||
|
|
|
||
|
|
// Regenerate the blob
|
||
|
|
out := bytes.NewBuffer(nil)
|
||
|
|
if err := printer.Fprint(out, fileSet, tree); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
blob = out.Bytes()
|
||
|
|
}
|
||
|
|
return blob, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Flattens a main file with rewritten dependencies.
|
||
|
|
func flatten(pkg string, main []byte, deps [][]byte, remDead bool) ([]byte, error) {
|
||
|
|
buffer := new(bytes.Buffer)
|
||
|
|
|
||
|
|
// Dump all code pieces into a buffer
|
||
|
|
fmt.Fprintf(buffer, "package %s\n\n", pkg)
|
||
|
|
fmt.Fprintf(buffer, "%s\n\n", main)
|
||
|
|
for _, dep := range deps {
|
||
|
|
fmt.Fprintf(buffer, "%s\n\n", dep)
|
||
|
|
}
|
||
|
|
// Format the blob to Go standards and add imports
|
||
|
|
blob, err := imports.Process("", buffer.Bytes(), nil)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Shake all dependencies, format and goimport again
|
||
|
|
if remDead {
|
||
|
|
if blob, err = shake(blob); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if blob, err = imports.Process("", blob, nil); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return blob, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Merges a Go source with all its dependencies.
|
||
|
|
func Merge(path string, shake bool) (string, error) {
|
||
|
|
// Collect the dependency chain
|
||
|
|
deps, err := dependencies(path)
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
// Rewrite all the dependencies
|
||
|
|
pieces := make([][]byte, 0, len(deps))
|
||
|
|
for pkg, sources := range deps {
|
||
|
|
// Collect all the declarations in need of rewriting
|
||
|
|
decls, err := declarations(sources)
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
// Rewrite each of them in each source file
|
||
|
|
for _, src := range sources {
|
||
|
|
blob, err := rewrite(src, pkg, decls)
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
pieces = append(pieces, blob)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Rewrite the main file itself, append all dependencies
|
||
|
|
main, err := rewrite(path, "", nil)
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
blob, err := flatten("main", main, pieces, shake)
|
||
|
|
if err != nil {
|
||
|
|
return "", err
|
||
|
|
}
|
||
|
|
return string(blob), nil
|
||
|
|
}
|