diff --git a/buildozer/test_common.sh b/buildozer/test_common.sh
index 21a0cd23c..1b7828235 100644
--- a/buildozer/test_common.sh
+++ b/buildozer/test_common.sh
@@ -15,7 +15,7 @@ function set_up() {
 # BUILD file contents ($1) to STDIN.
 function run() {
   options=(--buildifier=)
-  while [[ "$1" == --* ]]; 
+  while [[ "$1" == --* ]];
   do
     options+=("$1")
     shift
diff --git a/edit/edit.go b/edit/edit.go
index c47aed64d..4a50d3fc3 100644
--- a/edit/edit.go
+++ b/edit/edit.go
@@ -151,26 +151,32 @@ func PackageDeclaration(f *build.File) *build.Rule {
 	if pkg := ExistingPackageDeclaration(f); pkg != nil {
 		return pkg
 	}
-	all := []build.Expr{}
-	added := false
+	insertAfter := -1
 	call := &build.CallExpr{X: &build.Ident{Name: "package"}}
-	for _, stmt := range f.Stmt {
+	for i, stmt := range f.Stmt {
 		switch stmt.(type) {
 		case *build.CommentBlock, *build.LoadStmt, *build.StringExpr:
 			// Skip docstring, comments, and load statements to
 			// find a place to insert the package declaration.
-		default:
-			if !added {
-				all = append(all, call)
-				added = true
+			insertAfter = i
+			continue
+		case *build.CallExpr:
+			// Skip `workspace()` calls which have to be the very first statements
+			// of workspace files
+			if isWorkspaceCall(stmt) {
+				insertAfter = i
+				continue
 			}
+		default:
 		}
-		all = append(all, stmt)
-	}
-	if !added { // In case the file is empty.
-		all = append(all, call)
+		break
 	}
+	var all []build.Expr
+	all = append(all, f.Stmt[:insertAfter+1]...)
+	all = append(all, call)
+	all = append(all, f.Stmt[insertAfter+1:]...)
 	f.Stmt = all
+
 	return &build.Rule{call, ""}
 }
 
@@ -201,6 +207,15 @@ func isEmptyPackage(expr build.Expr) bool {
 	return false
 }
 
+func isWorkspaceCall(expr build.Expr) bool {
+	if call, ok := expr.(*build.CallExpr); ok {
+		if ident, ok := call.X.(*build.Ident); ok && ident.Name == "workspace" {
+			return true
+		}
+	}
+	return false
+}
+
 // InsertAfter inserts an expression after index i.
 func InsertAfter(i int, stmt []build.Expr, expr build.Expr) []build.Expr {
 	i = i + 1 // index after the element at i
@@ -1102,8 +1117,7 @@ func InsertLoad(stmts []build.Expr, location string, from, to []string) []build.
 
 	load := NewLoad(location, from, to)
 
-	var all []build.Expr
-	added := false
+	insertAfter := -1
 	for i, stmt := range stmts {
 		_, isComment := stmt.(*build.CommentBlock)
 		_, isString := stmt.(*build.StringExpr)
@@ -1114,32 +1128,26 @@ func InsertLoad(stmts []build.Expr, location string, from, to []string) []build.
 		// not valid.
 		// Pretend they're not there and skip past them while we look for
 		// possible workspace calls.
-		isSyntheticPackage := isEmptyPackage(stmt)
+		//isSyntheticPackage := isEmptyPackage(stmt)
 
 		// If we're editing a WORKSPACE file, bazel requires that the workspace
 		// declaration must be the very first expression in the WORKSPACE file,
 		// before any loads.
-		isWorkspaceCall := false
-		if callExpr, isCallExpr := stmt.(*build.CallExpr); isCallExpr {
-			if functionIdent, isIdent := callExpr.X.(*build.Ident); isIdent {
-				if functionIdent.Name == "workspace" {
-					isWorkspaceCall = true
-				}
-			}
-		}
+		isWorkspaceCallStmt := isWorkspaceCall(stmt)
 
-		if isComment || isDocString || isSyntheticPackage || isWorkspaceCall || added {
-			all = append(all, stmt)
+		if isComment || isDocString || isWorkspaceCallStmt {
+			insertAfter = i
 			continue
 		}
-		all = append(all, load)
-		all = append(all, stmt)
-		added = true
-	}
-	if !added { // Empty file or just comments.
-		all = append(all, load)
+		break
 	}
+
+	var all []build.Expr
+	all = append(all, stmts[:insertAfter+1]...)
+	all = append(all, load)
+	all = append(all, stmts[insertAfter+1:]...)
 	return all
+
 }
 
 // ReplaceLoad removes load statements for passed to-symbols and replaces them with a new