IdRefVisitor used in a neighbor provider

main.scala

//> using lib "software.amazon.smithy:smithy-model:1.26.4"
//> using lib "com.disneystreaming.smithy4s:smithy4s-protocol:0.17.1"
//> using scala "2.13.10"
import software.amazon.smithy.model.Model
import software.amazon.smithy.model.neighbor.Walker
import software.amazon.smithy.model.shapes.ShapeId
import scala.jdk.CollectionConverters._
import software.amazon.smithy.model.neighbor.NeighborProvider
import software.amazon.smithy.model.neighbor.Relationship
import software.amazon.smithy.model.shapes.Shape
import java.{util => ju}
import software.amazon.smithy.model.traits.Trait
import software.amazon.smithy.model.shapes.ShapeVisitor
import software.amazon.smithy.model.node.Node
import software.amazon.smithy.model.shapes.StringShape
import software.amazon.smithy.model.traits.IdRefTrait
import scala.jdk.OptionConverters._
import javax.management.relation.Relation
import software.amazon.smithy.model.neighbor.RelationshipType

object demo extends App {

  val model = Model
    .assembler()
    .discoverModels()
    .addUnparsedModel(
      "test.smithy",
      """$version: "2"
        |namespace test
        |
        |structure Hello {
        |  input: Input
        |}
        |
        |structure Input {
        | u: Union
        |}
        |
        |union Union {
        |  a: Thing
        |}
        |
        |@smithy4s.meta#adtMember(Union)
        |structure Thing {}
        |""".stripMargin,
    )
    .assemble()
    .unwrap()

  class IdRefVisitor(inShape: Shape, node: Node) extends ShapeVisitor.Default[List[Relationship]] {

    def getDefault(shape: Shape): List[Relationship] = Nil

    override def stringShape(
      shape: StringShape
    ): List[Relationship] = shape.getTrait(classOf[IdRefTrait]).asScala.toList.flatMap { idref =>
      idref.getSelector().select(model).asScala.toList.map { target =>
        // Q: what kind of relationship type should this be?
        Relationship.create(inShape, RelationshipType.MEMBER_TARGET, target)
      }
    }

  }

  def idrefProvider(model: Model, base: NeighborProvider): NeighborProvider =
    new NeighborProvider {

      def visitTrait(of: Shape, shapeId: ShapeId, trt: Trait): List[Relationship] =
        if (trt.getSourceLocation().getFilename().contains("test.smithy")) {
          model.expectShape(shapeId).accept(new IdRefVisitor(of, trt.toNode()))
        } else
          Nil

      def getNeighbors(shape: Shape): ju.List[Relationship] = {
        val extras = shape.getAllTraits().asScala.flatMap { case (k, v) => visitTrait(shape, k, v) }
        (base.getNeighbors(shape).asScala ++ extras).asJava
      }

    }

  new Walker(
    idrefProvider(model, NeighborProvider.withTraitRelationships(model, NeighborProvider.of(model)))
  )
    .walkShapes(model.expectShape(ShapeId.from("test#Thing")))
    .asScala
    .filterNot(_.getId().getNamespace().startsWith("smithy"))
    .foreach(println)

}