Implement diagram output

domainDocs4s-core/src/main/scala/domaindocs4s/collector/Collector.scala

@@ -3,12 +3,20 @@ package domaindocs4s.collector
 import domaindocs4s.errors.DomainDocsArgError
 import tastyquery.Contexts.Context
 import tastyquery.Symbols
-import tastyquery.Symbols.{DeclaringSymbol, Symbol, TermSymbol}
-import tastyquery.Trees._
+import tastyquery.Symbols.{ClassSymbol, DeclaringSymbol, Symbol, TermSymbol}
+import tastyquery.Trees.*
 
 import scala.collection.mutable.ListBuffer
 
-case class DocumentedSymbol(nameOverride: Option[String], description: Option[String], symbol: Symbol, path: Vector[DeclaringSymbol])
+case class DocumentedSymbol(
+    nameOverride: Option[String],
+    description: Option[String],
+    symbol: Symbol,
+    path: Vector[DeclaringSymbol],
+    declarations: Vector[TermSymbol],
+) {
+  def name: String = nameOverride.getOrElse(symbol.name.toString)
+}
 
 case class DocumentationTree(symbol: DocumentedSymbol, children: List[DocumentationTree])
 
@@ -66,6 +74,16 @@ class TastyQueryCollector(using ctx: Context) extends Collector {
       case _                              => tree
     }
 
+    def getDeclarations(symbol: Symbol): Vector[TermSymbol] = {
+      symbol match {
+        case sym: ClassSymbol =>
+          sym.declarations.collect {
+            case t: TermSymbol if !t.isSynthetic && t.name.toString != "<init>" => t
+          }.toVector
+        case _                => Vector()
+      }
+    }
+
     symbol match {
       case ts: TermSymbol if ts.isModuleVal =>
         None
@@ -86,6 +104,7 @@ class TastyQueryCollector(using ctx: Context) extends Collector {
               description = getConstArg(0, "description"),
               symbol = symbol,
               path = path,
+              declarations = getDeclarations(symbol),
             )
           })
     }

domainDocs4s-core/src/main/scala/domaindocs4s/collector/TastyContext.scala

@@ -4,22 +4,31 @@ import tastyquery.Contexts.Context
 import tastyquery.jdk.ClasspathLoaders
 
 import java.io.File
-import java.nio.file.Paths
+import java.net.URI
+import java.nio.file.{FileSystems, Path, Paths}
 
 type DomainDocsContext = Context
 
 object TastyContext {
 
   def fromCurrentProcess(): Context = {
-    val paths = sys
+    val appCp: List[Path] = sys
       .props("java.class.path")
       .split(File.pathSeparator)
       .toList
       .map(_.trim)
       .filter(_.nonEmpty)
-      .map(Paths.get(_).toAbsolutePath.normalize())
+      .map(p => Paths.get(p).toAbsolutePath.normalize())
 
-    val cp = ClasspathLoaders.read(paths)
+    val jdkBase: List[Path] =
+      try {
+        val jrtFs = FileSystems.getFileSystem(URI.create("jrt:/"))
+        List(jrtFs.getPath("modules", "java.base"))
+      } catch {
+        case _: Throwable => Nil
+      }
+
+    val cp = ClasspathLoaders.read(appCp ++ jdkBase)
     Context.initialize(cp)
   }
 

domainDocs4s-core/src/main/scala/domaindocs4s/output/Writer.scala

@@ -0,0 +1,12 @@
+package domaindocs4s.output
+
+import java.nio.file.{Files, Path}
+
+// TODO: for now it's only for decouple writing from generating, later it should be extended for more complex output handling
+object Writer {
+
+  def apply(docs: String, path: String): Unit = {
+    val _ = Files.write(Path.of(path), docs.getBytes)
+  }
+
+}

domainDocs4s-core/src/main/scala/domaindocs4s/output/diagram/Association.scala

@@ -0,0 +1,7 @@
+package domaindocs4s.output.diagram
+
+enum Association {
+  case ExactlyOne
+  case ZeroOrOne
+  case ZeroOrMore
+}

domainDocs4s-core/src/main/scala/domaindocs4s/output/diagram/Diagram.scala

@@ -0,0 +1,106 @@
+package domaindocs4s.output.diagram
+
+import domaindocs4s.collector.{Documentation, DocumentedSymbol}
+import tastyquery.Contexts.Context
+import tastyquery.Symbols.{Symbol, TermSymbol}
+import tastyquery.Trees.{DefDef, ValDef}
+import tastyquery.Types.{AppliedType, Type, TypeRef}
+
+case class Diagram(entities: List[Entity]) {
+
+  def asMarkdown(direction: Direction = Direction.None): String = {
+    if (entities.isEmpty) {
+      s"""
+         |```mermaid
+         |erDiagram
+         |```
+         |""".stripMargin
+    } else {
+      val directionLine = if (direction == Direction.None) "" else s"  direction $direction\n"
+      s"""
+         |```mermaid
+         |erDiagram
+         |$directionLine${entities.map(_.asMarkdown).mkString("\n")}
+         |```
+         |""".stripMargin
+    }
+  }
+
+}
+
+object Diagram {
+
+  def build(docs: Documentation)(using ctx: Context): Diagram = {
+    val documentedSymbols = docs.symbols
+    val docEntities       = docs.symbols.collect {
+      case sym if sym.symbol.isClass => buildDocEntity(sym, documentedSymbols)(using ctx)
+    }
+    Diagram(docEntities)
+  }
+
+  private def buildDocEntity(sym: DocumentedSymbol, documentedSymbols: List[DocumentedSymbol])(using ctx: Context): Entity = {
+    Entity(
+      name = sym.name,
+      associations = buildAssociations(sym, documentedSymbols)(using ctx),
+      fields = buildFields(sym),
+    )
+  }
+
+  private def buildAssociations(
+      sym: DocumentedSymbol,
+      documentedSymbols: List[DocumentedSymbol],
+  )(using ctx: Context): Map[String, Association] = {
+    sym.declarations.flatMap { field =>
+      documentedSymbols
+        .filter(s => isFieldRelatedToSymbol(field, s.symbol)(using ctx))
+        .map(s => (s.name, determineRelationshipType(field)(using ctx)))
+    }.toMap
+  }
+
+  private def buildFields(sym: DocumentedSymbol): Map[String, String] = {
+    sym.declarations.map(f => (f.name.toString, shortTypeName(typeOfTerm(f)))).toMap
+  }
+
+  private def isFieldRelatedToSymbol(field: TermSymbol, symbol: Symbol)(using ctx: Context): Boolean = {
+    termTypeSymbol(field)(using ctx).contains(symbol) || isTypeArgument(field, symbol)(using ctx)
+  }
+
+  private def isTypeArgument(field: TermSymbol, symbol: Symbol)(using ctx: Context): Boolean = {
+    def checkType(tp: Type): Boolean = tp match {
+      case at: AppliedType =>
+        at.args.exists {
+          case argType: Type => checkType(argType)
+          case _             => false
+        }
+      case tr: TypeRef     => tr.optSymbol.contains(symbol)
+      case _               => false
+    }
+
+    checkType(typeOfTerm(field))
+  }
+
+  private def determineRelationshipType(field: TermSymbol)(using ctx: Context): Association = {
+    val fieldType       = typeOfTerm(field).dealias
+    val optionCls       = ctx.findTopLevelClass("scala.Option")
+    val iterableOnceCls = ctx.findTopLevelClass("scala.collection.IterableOnce")
+
+    if (fieldType.baseType(optionCls).isDefined) Association.ZeroOrOne
+    else if (fieldType.baseType(iterableOnceCls).isDefined) Association.ZeroOrMore
+    else Association.ExactlyOne
+  }
+
+  private def typeOfTerm(t: TermSymbol): Type = t.tree match {
+    case Some(v: ValDef) => v.tpt.toType
+    case Some(d: DefDef) => d.resultTpt.toType
+    case Some(other)     => throw new Exception(s"Unexpected tree type for $t: ${other.getClass.getSimpleName}")
+    case None            => throw new Exception(s"No tree found for term symbol $t")
+  }
+
+  private def termTypeSymbol(t: TermSymbol)(using ctx: Context): Option[Symbol] = typeOfTerm(t).dealias match {
+    case tr: TypeRef => tr.optSymbol
+    case _           => None
+  }
+
+  private def shortTypeName(tp: Type): String =
+    tp.showBasic.replaceAll("\\b(?:[A-Za-z_][$\\w]*\\.)+([A-Za-z_][$\\w]*)\\b", "$1")
+}

domainDocs4s-core/src/main/scala/domaindocs4s/output/diagram/Direction.scala

@@ -0,0 +1,9 @@
+package domaindocs4s.output.diagram
+
+enum Direction {
+  case TB
+  case BT
+  case LR
+  case RL
+  case None
+}

domainDocs4s-core/src/main/scala/domaindocs4s/output/diagram/Entity.scala

@@ -0,0 +1,32 @@
+package domaindocs4s.output.diagram
+
+case class Entity(
+    name: String,
+    fields: Map[String, String],
+    associations: Map[String, Association],
+) {
+
+  def asMarkdown: String = {
+    val entityDefinition   = renderEntityDefinition
+    val entityAssociations = renderAssociations
+    escapeMermaidSpecialChars(entityDefinition + entityAssociations)
+  }
+
+  private def renderEntityDefinition: String =
+    if (fields.isEmpty) name
+    else name + fields.map((fieldName, fieldType) => s"$fieldName $fieldType").mkString(" {\n ", "\n ", "\n}")
+
+  private def renderAssociations: String =
+    associations.map((childName, relType) => s"\n$name ${associationAsMermaid(relType)} $childName : has").mkString
+
+  private def associationAsMermaid(rel: Association): String = rel match {
+    case Association.ExactlyOne => "||--||"
+    case Association.ZeroOrOne  => "||--o|"
+    case Association.ZeroOrMore => "||--o{"
+  }
+
+  private def escapeMermaidSpecialChars(text: String): String =
+    text
+      .replace("\"", "\\\"") // Escape double quotes
+      .replace("$", "")      // Remove dollar signs to prevent Mermaid parsing issues
+}

domainDocs4s-core/src/main/scala/domaindocs4s/output/Entry.scala → domainDocs4s-core/src/main/scala/domaindocs4s/output/glossary/Entry.scala

@@ -1,4 +1,4 @@
-package domaindocs4s.output
+package domaindocs4s.output.glossary
 
 import scala.math.Ordering.Implicits.given
 

domainDocs4s-core/src/main/scala/domaindocs4s/output/Glossary.scala → domainDocs4s-core/src/main/scala/domaindocs4s/output/glossary/Glossary.scala

@@ -1,9 +1,6 @@
-package domaindocs4s.output
+package domaindocs4s.output.glossary
 
 import domaindocs4s.collector.Documentation
-import domaindocs4s.output.Entry
-
-import java.nio.file.{Files, Path}
 
 case class Glossary(entries: List[Entry]) {
 
@@ -86,8 +83,4 @@ object Glossary {
 
     Glossary(cache.values.toList.sorted)
   }
-
-  def write(docs: String, path: String): Unit = {
-    val _ = Files.write(Path.of(path), docs.getBytes)
-  }
 }

domainDocs4s-examples/src/main/scala/domaindocs4s/banking/Main.scala

@@ -1,27 +1,52 @@
 package domaindocs4s.banking
 
-import domaindocs4s.collector.{DomainDocsContext, TastyContext, TastyQueryCollector}
-import domaindocs4s.output.Glossary
+object Main extends App {
 
-object Main {
+  // start_collector
+  import domaindocs4s.collector.{DomainDocsContext, TastyContext, TastyQueryCollector}
 
-  def main(args: Array[String]): Unit = {
+  // Setup collector
+  given DomainDocsContext = TastyContext.fromCurrentProcess()
+  val collector           = new TastyQueryCollector
 
-    // 1) Setup collector
-    given DomainDocsContext = TastyContext.fromCurrentProcess()
-    val collector           = new TastyQueryCollector
+  // Collect documentation from your domain package
+  val docs = collector
+    .collectSymbols("domaindocs4s.banking.application")
+  // end_collector
 
-    // 2) Collect documentation model
-    val docs = collector
-      .collectSymbols("domaindocs4s.banking")
+  // start_glossary
+  import domaindocs4s.output.glossary.Glossary
 
-    // 3) Generate glossary markdown (or another supported output and format)
-    val glossary =
-      Glossary
-        .build(docs) // build glossary
-        .asMarkdown  // render markdown
+  // Build glossary
+  val glossary = Glossary.build(docs)
+  // end_glossary
 
-    // 4) Write the result to file
-    Glossary.write(glossary, "glossary.md")
-  }
+  // start_md_glossary
+  import domaindocs4s.output.Writer
+
+  val glossaryMd = glossary.asMarkdown
+  Writer(glossaryMd, "glossary.md")
+  // end_md_glossary
+
+  // start_html_glossary
+  import domaindocs4s.output.Writer
+
+  val glossaryHtml = glossary.asHtml
+  Writer(glossaryHtml, "glossary.html")
+  // end_html_glossary
+
+  // start_diagram
+  // Build diagram and render it as markdown
+  import domaindocs4s.output.diagram.Diagram
+  import domaindocs4s.output.diagram.Direction
+  import domaindocs4s.output.Writer
+
+  val diagram =
+    Diagram
+      .build(docs)
+      .asMarkdown(Direction.TB)
+
+  // Write diagram to file
+  Writer(diagram, "diagram.md")
+  // end_diagram
 }

domainDocs4s-examples/src/main/scala/domaindocs4s/banking/application/Application.scala

@@ -7,30 +7,17 @@ import domaindocs4s.domainDoc
 import java.time.LocalDateTime
 import java.util.UUID
 
+// start_application
 @domainDoc("Loan application submitted by the customer for processing.")
 final case class Application(
     id: Application.Id,
     applicant: Party,
-    income: List[Income],
-    liability: List[Liability],
+    income: Option[Income],
+    liabilities: List[Liability],
     terms: Application.LoanTerms,
     status: ApplicationStatus,
     submittedAt: LocalDateTime,
-) {
-
-  @domainDoc("Marks the loan application as submitted by the customer.")
-  def submit(): Unit = ()
-
-  @domainDoc("Approves the loan application after successful review.")
-  def approve(): Unit = ()
-
-  @domainDoc("Rejects the loan application, preventing further processing.")
-  def reject(): Unit = ()
-
-  @domainDoc("Disburses the approved loan amount to the customer.")
-  def disburse(): Unit = ()
-
-}
+)
 
 object Application {
 
@@ -49,3 +36,4 @@ enum ApplicationStatus {
   case Rejected(reason: String)
   case Disbursed
 }
+// end_application