kickc/src/main/java/dk/camelot64/kickc/asm/AsmLibrary.java

package dk.camelot64.kickc.asm;

import dk.camelot64.kickc.model.CallGraph;
import dk.camelot64.kickc.model.Program;
import dk.camelot64.kickc.model.Registers;
import dk.camelot64.kickc.model.symbols.Procedure;
import dk.camelot64.kickc.model.symbols.Scope;
import dk.camelot64.kickc.model.symbols.Symbol;
import dk.camelot64.kickc.model.symbols.Variable;
import dk.camelot64.kickc.model.types.SymbolType;
import dk.camelot64.kickc.model.types.SymbolTypeStruct;
import dk.camelot64.kickc.model.values.*;

import java.nio.file.Path;
import java.util.*;

/**
 * A Kick assembler "library" helper class used to import and export assembler from a C source.
 */
public class AsmLibrary extends AsmLine {

   public String getAsmLibraryName() {
      return this.asmLibraryName.toString();
   }

   public String getAsmLibraryIdentifier() {
      return this.asmLibraryName.getAsm();
   }

   private final AsmIdentifier asmLibraryName;

   private final HashMap<String, Procedure.CallingConvention> symbols;

   public AsmLibrary(String asmLibraryName, Path resource, HashMap<String, Procedure.CallingConvention> symbols) {
      if(asmLibraryName!=null) {
         this.asmLibraryName = new AsmIdentifier(asmLibraryName);
      } else {
         this.asmLibraryName = null;
      }
      this.symbols = symbols;
   }

   public AsmLibrary(String asmLibraryName, Path resource) {
      this(asmLibraryName, resource, new LinkedHashMap<>());
   }

   public Set<String> getSymbols() {
      return this.symbols.keySet();
   }

   public void addSymbol(String symbolName, Procedure.CallingConvention callingConvention) {
      symbols.put(symbolName, callingConvention);
   }

   public Procedure.CallingConvention getProcedureCallingConvention(String procedureName) {
      return symbols.get(procedureName);
   }

   public boolean hasSymbol(String symbolName) {
      return symbols.containsKey(symbolName);
   }


   /**
    * Generate the zeropages used
    *
    * @param param     The variable to investigate
    * @param scope     The scope (procedure)
    * @param zp        The zero pages to append
    */
   void generateZPVar(Set<String> zp, Variable param, Scope scope) {
      if (param == null) return;
      if (param.isKindPhiMaster()) {
         List<Variable> versions = new ArrayList<>(scope.getVersions(param));
         if (!versions.isEmpty()) if (param.getLocalName().equals("return")) {
            // Choose the last version for return values
            param = versions.get(versions.size() - 1);
         } else {
            // Choose the first version for parameters
            param = versions.get(0);
         }
         else
            // Parameter optimized away to a constant or unused
            return;
      }

      Registers.Register allocation = param.getAllocation();
      if (allocation instanceof Registers.RegisterZpMem registerZp) {
         for(byte b=0; b<allocation.getBytes(); b++) {
            zp.add(Integer.toString(registerZp.getZp()+b));
         }
      }
   }

   void generatePrototypeVar(Variable var, StringBuilder signature) {
      Registers.Register allocation = var.getAllocation();
      if (allocation instanceof Registers.RegisterZpMem registerZp) {
          signature.append("__zp(").append(AsmFormat.getAsmNumber(registerZp.getZp())).append(") ");
      } else if (allocation instanceof Registers.RegisterMainMem registerMainMem) {
          signature.append("__mem(").append(registerMainMem.getAddress() == null ? "" : AsmFormat.getAsmNumber(registerMainMem.getAddress())).append(") ");
      } else if (allocation instanceof Registers.RegisterAByte) {
         signature.append("__register(A) ");
      } else if (allocation instanceof Registers.RegisterXByte) {
         signature.append("__register(X) ");
      } else if (allocation instanceof Registers.RegisterYByte) {
         signature.append("__register(Y) ");
      } else if (allocation instanceof Registers.RegisterZByte) {
         signature.append("__register(Z) ");
      }
   }

   /**
    * Generate part of a comment that describes a returnvalue/parameter
    *
    * @param param     The variable to describe
    * @param scope     The scope (procedure)
    * @param signature The signature to append to
    * @param isStackCall true if the signature is for stack call.
    * @return The version of the variable chosen
    */
   Variable generateSignatureVar(Variable param, Scope scope, StringBuilder signature, boolean isStackCall) {
      if (param == null) return param;
      if (param.isKindPhiMaster()) {
         List<Variable> versions = new ArrayList<>(scope.getVersions(param));
         if (versions.size() > 0) if (param.getLocalName().equals("return")) {
            // Choose the last version for return values
            param = versions.get(versions.size() - 1);
         } else {
            // Choose the first version for parameters
            param = versions.get(0);
         }
         else
            // Parameter optimized away to a constant or unused
            return param;
      }

      generatePrototypeVar(param, signature);
      return param;
   }

   /**
    * #820/40 - Reserve used zero-page registers usage by the procedure, if any.
    * If there are nested calls within the procedure to other exported procedures,
    * then the zero-page reserve must include all the zero-pages
    * of those nested functions!
    *
    * @param zp The set of consolidated zero-page used.
    * @param procedureRef The reference to the procedure to generate the nested zero-page usage for.
    * @param program The program.
    * @return The set of reserved zero-page collected in String type format.
    */
   private Set<String> generateZP(Set<String> zp, ProcedureRef procedureRef, Program program) {
      CallGraph callGraph = program.getCallGraph();
      Collection<ScopeRef> calls =
              callGraph.getRecursiveCalls(procedureRef);
      for(ScopeRef called : calls) {
         if(called instanceof ProcedureRef calledRef) {
            generateZP(zp, calledRef, program);
         }
      }
      Procedure procedure = program.getScope().getProcedure(procedureRef);
      HashMap<String, Symbol> procedureSymbols = procedure.getSymbols();
      for (Symbol symbol : procedureSymbols.values()) {
         if (symbol instanceof Variable variable) {
            generateZPVar(zp, variable, procedure);
         }
      }
      return zp;
   }

   /**
    * Generate a header for a C function prototype for .asm library import.
    * It defines the signature of each procedure in terms of the registers used,
    * which can be __mem, __zp, __register for the procedure parameters and return value.
    * Local variables, when allocated to __zp within the pre-compiled .asm library,
    * are reserved using __zp_reserve directive.
    * __stackcall procedures don't produce __mem, __zp, __register directives.
    * @param procedure The procedure to generate the header for.
    * @return A String with the generated function prototype.
    */
   private String generateProcedureHeader(Procedure procedure, Program program) {
      StringBuilder signature = new StringBuilder();
      Procedure.CallingConvention callingConvention = procedure.getCallingConvention();
      boolean isStackCall = callingConvention == Procedure.CallingConvention.STACK_CALL;
      signature.append("extern ");
      signature.append(callingConvention.getName()).append(" ");
      signature.append("__asm_import(\"").append(this.getAsmLibraryName()).append("\") ");

      Set<String> zpSet = generateZP(new HashSet<String>(), procedure.getRef(), program);
      if(!zpSet.isEmpty()) {
         String zpDelimited = String.join(",", zpSet);
         signature.append("__zp_reserve( ");
         signature.append(zpDelimited);
         signature.append(" ) ");
      }

      generateSignatureVar(procedure.getLocalVar("return"), procedure, signature, isStackCall);
      signature.append(procedure.getReturnType().toCDecl());
      signature.append(" ").append(procedure.getLocalName()).append("(");
      int i = 0;
      for (Variable parameter : procedure.getParameters()) {
         if (i++ > 0) signature.append(", ");
         Variable param = generateSignatureVar(parameter, procedure, signature, isStackCall);
         signature.append(param.getType().toCDecl(parameter.getLocalName()));
      }
      signature.append(");");
      signature.append("\n");
      // Always add the signature comments...
      return signature.toString();
   }

   /**
    * Generate struct forward declarations for C function prototypes for .asm library import,
    * that use structs in the parameters or return values.
    * @param procedure The procedure to generate the struct forward for.
    * @return A String with the generated function prototype.
    */
   private String generateProcedureStructForwards(Procedure procedure, Program program) {
      StringBuilder signature = new StringBuilder();
      HashSet<String> structForwards = new HashSet<>();

      Variable returnVariable = procedure.getLocalVar("return");
      if(returnVariable != null && returnVariable.isStruct()) {
         String structReturn = generateVariableStructForward(returnVariable, program);
         structForwards.add(structReturn);
      }

      for (Variable parameter : procedure.getParameters()) {
         if(parameter.isStruct()) {
            String paramStruct = generateVariableStructForward(parameter, program);
            structForwards.add(paramStruct);
         }
      }
      if(!structForwards.isEmpty())
         signature.append(String.join(";\n", structForwards)).append(";\n");
      return signature.toString();
   }

   private boolean hasData(Variable constantVar) {
      ConstantValue constantValue = constantVar.getInitValue();
      if (constantValue instanceof ConstantArray) return true;
      else if (constantValue instanceof ConstantStructValue) return true;
      else if (constantValue instanceof ConstantString) return true;
      else return false;
   }

   private boolean hasExportAsmLibrary(Variable constantVar, String exportAsmLibrary) {
      String varExportAsmLibrary = constantVar.getExportAsmLibrary();
      if(exportAsmLibrary != null) {
         if(varExportAsmLibrary != null) {
            return constantVar.getExportAsmLibrary().equals(exportAsmLibrary);
         } else {
            return false;
         }
      } else {
         return true;
      }
   }

   private String generateVariableStructForward(Variable variable, Program program) {
      StringBuilder signature = new StringBuilder();
      // #820/50 - Handle struct definitions correctly for proper declaration.
      if(variable.isStruct()) {
         SymbolType symbolType = variable.getType();
         if(symbolType instanceof SymbolTypeStruct typeStruct) {
            if(typeStruct.getTypeDefName() == null) {
               String structName = variable.getType().toCDecl();
               signature.append(structName).append(";\n");
            }
         }
      }
      return signature.toString();

   }

   private String generateVariableHeader(Variable variable, Program program) {
      StringBuilder signature = new StringBuilder();
      signature.append("extern ");
      signature.append("__asm_import(\"").append(this.getAsmLibraryName()).append("\") ");
      generatePrototypeVar(variable, signature);
//      signature.append(variable.getType().toCDecl());
//      signature.append(variable.getLocalName()).append(" ");
      signature.append(variable.toCDecl());
//      if(variable.isKindConstant()) {
//         signature.append(" = ").append(variable.getInitValue());
//      }
      signature.append(";");
      signature.append("\n");
      // Always add the signature comments...
      return signature.toString();

   }
      /**
       * Generate a header file for an .asm library import.
       * This is to be used by the main program to import an .asm library.
       * It defines for each the registers used, which can be `__mem`, `__zp`, `__register`:
       * - the extern declaration of each exported global variable in the .asm library.
       * - the extern declaration of each exported procedure in the .asm library,
       *   in terms of the registers used, for the procedure parameters and return value.
       *   Local variables, when allocated to `__zp` within the pre-compiled .asm library,
       *   are reserved using `__zp_reserve` directive.
       * @param program The program
       * @return A string with the generated function prototypes for import.
       */
   public String generateHeaders(Program program) {
      StringBuilder headers = new StringBuilder();
      StringBuilder structs = new StringBuilder();

      // Generate the global variables headers.
      Scope scope  = program.getScope().getScope(ScopeRef.ROOT);
      Collection<Variable> globalVariables = scope.getAllVars(false);
      Set<String> zp = new HashSet<>();
      // Add all constants arrays incl. strings with data
      for (Variable globalVariable : globalVariables) {
         if (hasExportAsmLibrary(globalVariable, this.getAsmLibraryName())) {
            if(globalVariable.isMemoryAreaZP()) {
               Registers.RegisterZpMem registerZpMem = (Registers.RegisterZpMem)globalVariable.getAllocation();
               if(registerZpMem != null) {
                   for(byte b=0; b<registerZpMem.getBytes(); b++) {
                       zp.add(Integer.toString(registerZpMem.getZp()+b));
                   }
               }
            }
            // TODO: #820/50 - Declaration of global variables only when they are exported using asm_export pragma.
            // This requires rework. It is a challenge to generate only those global var
            // By default at this moment, all global variables in libraries are encapsulated in the asm namespace.
            // This prevents any conflict of global variables.
             headers.append(generateVariableHeader(globalVariable, program));
            structs.append(generateVariableStructForward(globalVariable, program));
         }
      }

      if(!zp.isEmpty())
         headers.insert(0, "#pragma zp_reserve(" + String.join(",", zp) + ")\n");

      // Generate the procedure headers.
      for(String procedureName: this.getSymbols()) {
         Procedure procedure = program.getScope().getProcedure(new ProcedureRef(procedureName));
         if(procedure != null) {
            if(procedure.isAsmExportLibrary())
               headers.append(generateProcedureHeader(procedure, program));
               structs.append(generateProcedureStructForwards(procedure, program));
         }
      }
      return structs.toString() + headers.toString();
   }

   @Override
   public int getLineBytes() {
      return 0;
   }

   @Override
   public double getLineCycles() {
      return 0;
   }

   @Override
   public String getAsm() {
      StringBuilder asm = new StringBuilder();
      asm.append(".export ").append(" [ ");
      asm.append("name=\"").append(asmLibraryName).append("\"");
      for(String procedureName : symbols.keySet()) {
         asm.append(",");
         asm.append(procedureName);
      }
      asm.append(" ]");
      return asm.toString();
   }

}