Add variable content to PageNodes and extras to DocumentationNodes

I'm just sorry

1 files changed, 38 insertions, 6 deletions

diff --git a/core/src/main/kotlin/Model/DocumentationNode.kt b/core/src/main/kotlin/Model/DocumentationNode.kt
index fc81db6d..55ae5902 100644
--- a/core/src/main/kotlin/Model/DocumentationNode.kt
+++ b/core/src/main/kotlin/Model/DocumentationNode.kt
@@ -8,13 +8,15 @@ import org.jetbrains.kotlin.kdoc.psi.impl.KDocTag
 class Module(val packages: List<Package>) : DocumentationNode<Nothing>() {
     override val dri: DRI = DRI.topLevel
     override val children: List<Package> = packages
+    override val extra: MutableSet<Extra> = mutableSetOf()
 }
 
 class Package(
     override val dri: DRI,
     override val functions: List<Function>,
     override val properties: List<Property>,
-    override val classes: List<Class>
+    override val classes: List<Class>,
+    override val extra: MutableSet<Extra> = mutableSetOf()
 ) : ScopeNode<Nothing>() {
     val name = dri.packageName.orEmpty()
 }
@@ -27,7 +29,8 @@ class Class(
     override val properties: List<Property>,
     override val classes: List<Class>,
     override val expectDescriptor: Descriptor<ClassDescriptor>?,
-    override val actualDescriptors: List<Descriptor<ClassDescriptor>>
+    override val actualDescriptors: List<Descriptor<ClassDescriptor>>,
+    override val extra: MutableSet<Extra> = mutableSetOf()
 ) : ScopeNode<ClassDescriptor>()
 
 class Function(
@@ -36,7 +39,8 @@ class Function(
     override val receiver: Parameter?,
     val parameters: List<Parameter>,
     override val expectDescriptor: Descriptor<FunctionDescriptor>?,
-    override val actualDescriptors: List<Descriptor<FunctionDescriptor>>
+    override val actualDescriptors: List<Descriptor<FunctionDescriptor>>,
+    override val extra: MutableSet<Extra> = mutableSetOf()
 ) : CallableNode<FunctionDescriptor>() {
     override val children: List<Parameter>
         get() = listOfNotNull(receiver) + parameters
@@ -47,7 +51,8 @@ class Property(
     val name: String,
     override val receiver: Parameter?,
     override val expectDescriptor: Descriptor<PropertyDescriptor>?,
-    override val actualDescriptors: List<Descriptor<PropertyDescriptor>>
+    override val actualDescriptors: List<Descriptor<PropertyDescriptor>>,
+    override val extra: MutableSet<Extra> = mutableSetOf()
 ) : CallableNode<PropertyDescriptor>() {
     override val children: List<Parameter>
         get() = listOfNotNull(receiver)
@@ -57,7 +62,8 @@ class Property(
 class Parameter(
     override val dri: DRI,
     val name: String?,
-    override val actualDescriptors: List<Descriptor<ParameterDescriptor>>
+    override val actualDescriptors: List<Descriptor<ParameterDescriptor>>,
+    override val extra: MutableSet<Extra> = mutableSetOf()
 ) : DocumentationNode<ParameterDescriptor>() {
     override val children: List<DocumentationNode<*>>
         get() = emptyList()
@@ -103,6 +109,8 @@ abstract class DocumentationNode<out T : DeclarationDescriptor> {
 
     val briefDocstring: String
         get() = descriptors.firstOrNull()?.docTag?.getContent().orEmpty().shorten(40)
+
+    open val extra: MutableSet<Extra> = mutableSetOf()
 }
 
 abstract class ScopeNode<out T : ClassOrPackageFragmentDescriptor> : DocumentationNode<T>() {
@@ -120,4 +128,28 @@ abstract class CallableNode<out T : CallableDescriptor> : DocumentationNode<T>()
 
 private fun String.shorten(maxLength: Int) = lineSequence().first().let {
     if (it.length != length || it.length > maxLength) it.take(maxLength - 3) + "..." else it
-}
\ No newline at end of file
+}
+
+fun DocumentationNode<*>.walk(process: DocumentationNode<*>.() -> Unit) {
+    this.process()
+    this.children.forEach { it.process() }
+}
+
+fun DocumentationNode<*>.dfs(predicate: (DocumentationNode<*>) -> Boolean): DocumentationNode<*>? =
+    if (predicate(this)) {
+        this
+    } else {
+        this.children.asSequence().mapNotNull { it.dfs(predicate) }.firstOrNull()
+    }
+
+fun DocumentationNode<*>.findAll(predicate: (DocumentationNode<*>) -> Boolean): Set<DocumentationNode<*>> {
+    val found = mutableSetOf<DocumentationNode<*>>()
+    if (predicate(this)) {
+        found.add(this)
+    } else {
+        this.children.asSequence().mapNotNull { it.findAll(predicate) }.forEach { found.addAll(it) }
+    }
+    return found
+}
+
+interface Extra
\ No newline at end of file