Better entity code generation + minor refactoring

annotations/src/main/java/mindustry/annotations/BaseProcessor.java

@@ -10,6 +10,8 @@ import javax.lang.model.*;
 import javax.lang.model.element.*;
 import javax.lang.model.util.*;
 import javax.tools.Diagnostic.*;
+import javax.tools.*;
+import java.io.*;
 import java.lang.annotation.*;
 import java.util.*;
 
@@ -38,7 +40,33 @@ public abstract class BaseProcessor extends AbstractProcessor{
     }
 
     public static void write(TypeSpec.Builder builder) throws Exception{
-        JavaFile.builder(packageName, builder.build()).build().writeTo(BaseProcessor.filer);
+        write(builder, null);
+    }
+
+    public static void write(TypeSpec.Builder builder, Array<String> imports) throws Exception{
+        JavaFile file = JavaFile.builder(packageName, builder.build()).skipJavaLangImports(true).build();
+
+        if(imports != null){
+            String rawSource = file.toString();
+            Array<String> result = new Array<>();
+            for (String s : rawSource.split("\n", -1)) {
+                result.add(s);
+                if (s.startsWith("package ")) {
+                    result.add("");
+                    for (String i : imports) {
+                        result.add(i);
+                    }
+                }
+            }
+
+            String out = result.toString("\n");
+            JavaFileObject object = filer.createSourceFile(file.packageName + "." + file.typeSpec.name, file.typeSpec.originatingElements.toArray(new Element[0]));
+            OutputStream stream = object.openOutputStream();
+            stream.write(out.getBytes());
+            stream.close();
+        }else{
+            file.writeTo(filer);
+        }
     }
 
     public Array<Stype> types(Class<? extends Annotation> type){

annotations/src/main/java/mindustry/annotations/impl/CallSuperProcess.java

@@ -122,6 +122,7 @@ public class CallSuperProcess extends AbstractProcessor{
             }
 
             for(Symbol s : it){
+
                 if(s instanceof MethodSymbol){
                     MethodSymbol ms = (MethodSymbol)s;
 

annotations/src/main/java/mindustry/annotations/impl/EntityProcess.java

@@ -23,6 +23,7 @@ public class EntityProcess extends BaseProcessor{
     Array<Stype> baseComponents;
     ObjectMap<Stype, Array<Stype>> componentDependencies = new ObjectMap<>();
     ObjectMap<Stype, Array<Stype>> defComponents = new ObjectMap<>();
+    ObjectSet<String> imports = new ObjectSet<>();
 
     {
         rounds = 2;
@@ -40,6 +41,7 @@ public class EntityProcess extends BaseProcessor{
 
             //find all components used...
             for(Stype type : allDefs){
+                imports.addAll(getImports(type.e));
                 allComponents.addAll(allComponents(type));
             }
 
@@ -70,6 +72,7 @@ public class EntityProcess extends BaseProcessor{
                 //add utility methods to interface
                 for(Smethod method : component.methods()){
                     inter.addMethod(MethodSpec.methodBuilder(method.name())
+                    .addExceptions(method.thrownt())
                     .addTypeVariables(method.typeVariables().map(TypeVariableName::get))
                     .returns(method.ret().toString().equals("void") ? TypeName.VOID : method.retn())
                     .addParameters(method.params().map(v -> ParameterSpec.builder(v.tname(), v.name())
@@ -119,6 +122,7 @@ public class EntityProcess extends BaseProcessor{
                     MethodSpec.Builder mbuilder = MethodSpec.methodBuilder(first.name()).addModifiers(Modifier.PUBLIC, Modifier.FINAL);
                     mbuilder.addTypeVariables(first.typeVariables().map(TypeVariableName::get));
                     mbuilder.returns(first.retn());
+                    mbuilder.addExceptions(first.thrownt());
 
                     for(Svar var : first.params()){
                         mbuilder.addParameter(var.tname(), var.name());
@@ -127,19 +131,25 @@ public class EntityProcess extends BaseProcessor{
                     boolean returns = !first.ret().toString().equals("void");
 
                     for(Smethod elem : entry.value){
-                        //wrap scope to prevent variable leakage
-                        if(!returns) mbuilder.beginControlFlow("");
-
                         //get all statements in the method, copy them over
                         MethodTree methodTree = elem.tree();
                         BlockTree blockTree = methodTree.getBody();
-                        for(StatementTree st : blockTree.getStatements()){
-                            String state = st.toString();
-                            mbuilder.addStatement(state.substring(0, state.length() - 1));
+                        String str = blockTree.toString();
+                        String blockName = elem.up().getSimpleName().toString().toLowerCase().replace("comp", "");
+
+                        //skip empty blocks
+                        if(str.replace("{", "").replace("\n", "").replace("}", "").replace("\t", "").replace(" ", "").isEmpty()){
+                            continue;
                         }
 
+                        //wrap scope to prevent variable leakage
+                        if(!returns) mbuilder.addCode(blockName + ": {\n");
+
+                        //make sure to remove braces here
+                        mbuilder.addCode(str.substring(2, str.length() - 1).replace("return;", "break " + blockName + ";"));
+
                         //end scope
-                        if(!returns) mbuilder.endControlFlow();
+                        if(!returns) mbuilder.addCode("}\n");
                     }
 
                     builder.addMethod(mbuilder.build());
@@ -167,6 +177,9 @@ public class EntityProcess extends BaseProcessor{
                         if(method.name().length() <= 3) continue;
 
                         String var = Strings.camelize(method.name().substring(3));
+                        //make sure it's a real variable
+                        if(!Array.with(def.builder.fieldSpecs).contains(f -> f.name.equals(var))) continue;
+
                         if(method.name().startsWith("get")){
                             def.builder.addMethod(MethodSpec.overriding(method.e).addStatement("return " + var).build());
                         }else if(method.name().startsWith("set")){
@@ -175,11 +188,16 @@ public class EntityProcess extends BaseProcessor{
                     }
                 }
 
-                write(def.builder);
+                write(def.builder, imports.asArray());
             }
         }
     }
 
+    Array<String> getImports(Element elem){
+        return Array.with(trees.getPath(elem).getCompilationUnit().getImports()).map(t -> t.toString());
+    }
+
+    /** @return interface for a component type */
     String interfaceName(Stype comp){
         String suffix = "Comp";
         if(!comp.name().endsWith(suffix)){

annotations/src/main/java/mindustry/annotations/util/Selement.java

@@ -13,6 +13,10 @@ public class Selement<T extends Element>{
         this.e = e;
     }
 
+    public Element up(){
+        return e.getEnclosingElement();
+    }
+
     public TypeMirror mirror(){
         return e.asType();
     }

annotations/src/main/java/mindustry/annotations/util/Smethod.java

@@ -14,6 +14,14 @@ public class Smethod extends Selement<ExecutableElement>{
         super(executableElement);
     }
 
+    public Array<TypeMirror> thrown(){
+        return Array.with(e.getThrownTypes()).as(TypeMirror.class);
+    }
+
+    public Array<TypeName> thrownt(){
+        return Array.with(e.getThrownTypes()).map(TypeName::get);
+    }
+
     public Array<TypeParameterElement> typeVariables(){
         return Array.with(e.getTypeParameters()).as(TypeParameterElement.class);
     }