Example #1
0
// Track imports simplifies all imports into a single large package
// with mangled names.  In the process, it drops functions that are
// never referred to.
func TrackImports(pkgs map[string](map[string]*ast.File)) (main *ast.File) {
	// Let's first set of the package we're going to generate...
	main = new(ast.File)
	main.Name = ast.NewIdent("main")
	initstmts := []ast.Stmt{} // this is where we'll stash the init statements...

	todo := make(map[string]struct{})
	todo["main.init"] = struct{}{}
	todo["main.main"] = struct{}{}
	done := make(map[string]struct{})

	for len(todo) > 0 {
		for pkgfn := range todo {
			pkg := splitLast(pkgfn, ".")[0] // FIXME:  Need to split after last "." only
			fn := splitLast(pkgfn, ".")[1]
			fmt.Println("Working on", fn, "in", pkg)
			if _, ok := done[pkg+".init"]; !ok && fn != "init" {
				// We still need to init this package!
				todo[pkg+".init"] = struct{}{}
			}
			// fmt.Println("Working on package", pkg, "function", fn)
			// We need to look in all this package's files...
			for _, f := range pkgs[pkg] {
				// FIXME: it'd be marginally faster to first check if the
				// function we want is in this particular file.  On the other
				// hand, when there's only one file per package, that would be
				// slower...

				// First we'll track down the import declarations...
				sc := PackageScoping{
					Imports: make(map[string]string),
					Globals: make(map[string]string),
				}
				for _, d := range f.Decls {
					if i, ok := d.(*ast.GenDecl); ok && i.Tok == token.IMPORT {
						for _, s := range i.Specs {
							ispec := s.(*ast.ImportSpec) // This should always be okay!
							path, _ := strconv.Unquote(string(ispec.Path.Value))
							name := path
							if ispec.Name != nil {
								name = ispec.Name.Name
							} else {
								for _, f := range pkgs[path] {
									name = f.Name.Name
								}
							}
							sc.Imports[name] = path
						}
					} else if vdecl, ok := d.(*ast.GenDecl); ok && vdecl.Tok == token.VAR {
						for _, spec0 := range vdecl.Specs {
							spec := spec0.(*ast.ValueSpec)
							for _, n := range spec.Names {
								sc.Globals[n.Name] = pkg
							}
						}
					} else if tdecl, ok := d.(*ast.GenDecl); ok && vdecl.Tok == token.TYPE {
						for _, spec0 := range tdecl.Specs {
							spec := spec0.(*ast.TypeSpec)
							sc.Globals[spec.Name.Name] = pkg
						}
					} else if fdecl, ok := d.(*ast.FuncDecl); ok {
						sc.Globals[fdecl.Name.Name] = pkg
					}
				}
				// Now we'll go ahead and mangle things...
				for _, d := range f.Decls {
					if cdecl, ok := d.(*ast.GenDecl); ok && cdecl.Tok == token.CONST {
						fmt.Println("FIXME: I don't handle const yet at all... (ignoring)")
					} else if tdecl, ok := d.(*ast.GenDecl); ok && tdecl.Tok == token.TYPE {
						for _, spec0 := range tdecl.Specs {
							spec := spec0.(*ast.TypeSpec)
							if spec.Name.Name == fn {
								// fmt.Println("Got type declaration of", spec.Name)
								spec := *spec
								spec.Name = ast.NewIdent(fn)
								spec.Type = sc.MangleExpr(spec.Type)
								sc.MangleExpr(spec.Name)
								d := &ast.GenDecl{
									Tok:   token.TYPE,
									Specs: []ast.Spec{&spec},
								}
								main.Decls = append(main.Decls, d)
							}
						}
					} else if vdecl, ok := d.(*ast.GenDecl); ok && vdecl.Tok == token.VAR {
						for _, spec0 := range vdecl.Specs {
							spec := spec0.(*ast.ValueSpec)
							for i, n := range spec.Names {
								if n.Name == fn {
									// fmt.Println("I got variable", fn)
									nnew := *n
									sc.MangleExpr(&nnew)
									vs := []ast.Expr(nil)
									if len(spec.Values) > i {
										vs = append(vs, spec.Values[i])
										sc.MangleExpr(spec.Values[i])
									}
									sc.MangleExpr(spec.Type)
									d := ast.GenDecl{
										Tok: token.VAR,
										Specs: []ast.Spec{
											&ast.ValueSpec{
												Names:  []*ast.Ident{&nnew},
												Type:   spec.Type,
												Values: vs,
											},
										},
									}
									main.Decls = append(main.Decls, &d)
								}
							}
						}
					} else if fdecl, ok := d.(*ast.FuncDecl); ok {
						if fdecl.Name.Name == fn {
							// first, let's update the name... but in a copy of the
							// function declaration
							fdecl := *fdecl
							fdecl.Name = ast.NewIdent(pkg + "_" + fn)
							if fdecl.Type.Params != nil {
								for _, f := range fdecl.Type.Params.List {
									sc.MangleExpr(f.Type)
								}
							}
							if fdecl.Type.Results != nil {
								for _, f := range fdecl.Type.Results.List {
									sc.MangleExpr(f.Type)
								}
							}
							sc.MangleStatement(fdecl.Body)
							// fmt.Println("Dumping out", pkg, fn)
							main.Decls = append(main.Decls, &fdecl)
							if fn == "init" && fdecl.Recv == nil {
								initstmts = append(initstmts,
									&ast.ExprStmt{&ast.CallExpr{Fun: fdecl.Name}})
							}
						}
					}
				}
				// See what else we need to compile...
				for _, x := range sc.ToDo {
					if _, done := done[x]; !done {
						todo[x] = struct{}{}
					}
				}
			}
			delete(todo, pkgfn)
			done[pkgfn] = struct{}{}
		}
	}

	// Now we reverse the order, so that the declarations will be in
	// C-style order, not requiring forward declarations.
	newdecls := make([]ast.Decl, len(main.Decls))
	for i := range main.Decls {
		newdecls[i] = main.Decls[len(main.Decls)-i-1]
	}
	main.Decls = newdecls

	mainfn := new(ast.FuncDecl)
	mainfn.Name = ast.NewIdent("main")
	mainfn.Type = &ast.FuncType{Params: &ast.FieldList{}}
	initstmts = append(initstmts,
		&ast.ExprStmt{&ast.CallExpr{Fun: ast.NewIdent("main_main")}})
	mainfn.Body = &ast.BlockStmt{List: initstmts}
	main.Decls = append(main.Decls, mainfn)
	return main
}