LeanProtocPlugin.lean

import LeanProtocPlugin.Google.Protobuf.Compiler.Plugin
import LeanProtocPlugin.Google.Protobuf.Descriptor

import LeanProto
import LeanProtocPlugin.Helpers
import LeanProtocPlugin.ProtoGenM
import LeanProtocPlugin.Logic

import Lean.Elab.Term
import Lean.Hygiene
import Init.System.IO
import Std.Data.RBTree
import Std.Data.RBMap

open LeanProtocPlugin.Google.Protobuf.Compiler
open LeanProtocPlugin.Google.Protobuf
open IO (FS.Stream)

open Lean
open Lean.Elab.Term

-- deriving instance Inhabited for Std.RBMap

def defaultImports (sourceFilename: String) : Array String := #[
  "-- Generated by the Lean protobuf compiler. Do not edit manually.",
  s!"-- source: {sourceFilename}",
  "",
  "import LeanProto",
  "import Std.Data.AssocList"
]

def defaultSettings : Array String := #[
  "set_option maxHeartbeats 10000000",
  "set_option maxRecDepth 2048",
  "set_option synthInstance.maxHeartbeats 10000000",
  "set_option genSizeOfSpec false",
  "",
  "open Std (AssocList)"
]

def addImportDeps (f: FileDescriptorProto): ProtoGenM Unit := do
  addLine ""
  let namespacePrefix := (← read).namespacePrefix
  let paths := f.dependency.map (fun s => namespacePrefix ++ "." ++ filePathToPackage s)
  let importCommands ← paths.mapM (s!"import {·}")
  addLines importCommands

-- Building the context: the incoming proto request doesn't contain
-- enough information to do generation in a streaming way, so first
-- we go through everything and build up a context
namespace ContextPrep
  structure ContextPrepState where
    enums : Array (String × (ASTPath × EnumDescriptorProto)) := #[]
    oneofs : Array (String × (ASTPath × OneofDescriptorProto)) := #[]
    messages : Array (String × ASTPath) := #[]
    deriving Inhabited

  def prepareContextRecurseFile (d: ASTPath) (_: Unit): StateM ContextPrepState Unit := do
    let mut s ← get
    let extraS : ContextPrepState := match d.revMessages.head? with 
    | some m => {
      enums := m.enumType.map (fun e => (d.protoFullPath ++ "." ++ e.name, (d, e))),
      oneofs := m.oneofDecl.map (fun e => (d.protoFullPath ++ "." ++ e.name, (d, e))),
      messages := #[(d.protoFullPath, d)]
      }
    | none => {
      enums := d.file.enumType.map (fun e => (d.protoFullPath ++ "." ++ e.name, (d, e))),
    }

    set { 
      enums := s.enums ++ extraS.enums, 
      oneofs := s.oneofs ++ extraS.oneofs, 
      messages := s.messages ++ extraS.messages
      : ContextPrepState}

  def prepareContext (r: Compiler.CodeGeneratorRequest) (rootPackage: String): ContextPrepState :=   
    let rec doWork (f: FileDescriptorProto) : StateM ContextPrepState PUnit := do
      recurseM ((ASTPath.init f rootPackage), ()) (wrapRecurseFn prepareContextRecurseFile true)
    let doWorkM := r.protoFile.forM $ doWork
    StateT.run doWorkM {} |> Id.run |> Prod.snd
end ContextPrep

-- Generate code for a single .proto file
def doWork (file: FileDescriptorProto) : ProtoGenM $ List CodeGeneratorResponse_File := do
  addLines $ defaultImports file.name
  addImportDeps file
  addLine ""
  addLines $ defaultSettings
  addLine ""
    
  let package := (← read).namespacePrefix ++ "." ++ protoPackageToLeanPackagePrefix file.package
  addLine s!"namespace {package}"
  addLine ""

  generateEnumDeclarations file

  addLines #["", "mutual", ""]
  generateMessageDeclarations file
  addLines #["", "end", ""]

  generateMessageManualDerives file
  addLine ""
  generateLeanDeriveStatements file

  addLines #["", "mutual", ""]
  generateDeserializers file
  addLines #["", "end", ""]

  -- addLine s!"set_option trace.compiler.ir.result true"
  addLine s!""
  addLines #["", "mutual", ""]
  generateSerializers file
  addLines #["", "end", ""]
  -- addLine s!"set_option trace.compiler.ir.result false"
  addLine s!""

  generateSerDeserInstances file

  addLine ""
  addLine s!"end {package}"

  let protoFile := CodeGeneratorResponse_File.mk
    (outputFilePath file (← read).namespacePrefix) "" ("\n".intercalate (← get).lines.toList) none

  let services := file.service
  if services.size == 0 then return [protoFile]

  -- Reset line state
  set { (← get) with lines := #[] }
  
  addLines $ defaultImports file.name
  addImportDeps file

  let protoMod := (← read).namespacePrefix ++ "." ++ filePathToPackage file.name
  addLine s!"import LeanGRPC"
  addLine s!"import {protoMod}"
  addLine ""
  addLine s!"namespace {package}"
  addLine ""
  generateGRPC file
  addLine ""
  addLine s!"end {package}"
  let grpcFile := CodeGeneratorResponse_File.mk 
    (grpcOutputFilePath file (← read).namespacePrefix) "" ("\n".intercalate (← get).lines.toList) none
  return [protoFile, grpcFile]


def main(argv: List String): IO UInt32 := do
  let i ← IO.getStdin
  let o ← IO.getStdout

  let makeErrorResponse (s: String) := CodeGeneratorResponse.mk s 0 #[]  

  let bytes ← i.readBinToEnd
  let request ← monadLift $ LeanProto.ProtoDeserialize.deserialize (α:=CodeGeneratorRequest) bytes

  let namespacePrefix ← match request.parameter with
    | "" => do 
      let out ← LeanProto.ProtoSerialize.serialize $ makeErrorResponse "Must provide root package name as protoc parameter"
      o.write out
      return 0
    | _ => request.parameter

  let filesToWrite := rbtreeOfC request.fileToGenerate
  let fileProtos := rbmapOfC (request.protoFile.map fun x => (x.name, x))
  let fileProtosToWrite := request.protoFile.filter fun v => filesToWrite.contains v.name

  let contextAux := ContextPrep.prepareContext request namespacePrefix
  let ctxForFile (f: FileDescriptorProto) := ProtoGenContext.mk namespacePrefix fileProtos f
    (rbmapOfC contextAux.enums)
    (rbmapOfC contextAux.oneofs)
    (rbmapOfC contextAux.messages)

  let processFile (f: FileDescriptorProto) : IO $ List CodeGeneratorResponse_File :=
    StateT.run' (ReaderT.run (doWork f) (ctxForFile f)) {}  

  let response : CodeGeneratorResponse ← do try
    let processedNested : Array $ List CodeGeneratorResponse_File ← fileProtosToWrite.mapM processFile
    let mut processed := processedNested.foldl (fun acc curr => acc ++ curr.toArray) #[]

    let allNewModules := processed.foldl (fun acc s => (filePathToPackage s.name) :: acc) []
    let rootFileText := "\n".intercalate $ allNewModules.map (s!"import {·}")  
    processed := processed.push $ CodeGeneratorResponse_File.mk (namespacePrefix ++ ".lean") "" rootFileText none

    CodeGeneratorResponse.mk "" 0 processed
  catch e: IO.Error =>
    makeErrorResponse e.toString

  let r ← LeanProto.ProtoSerialize.serialize response
  o.write r
  return 0