Trait for composite primary keys in Slick tables?

Question

I have 2 tables

case class SimplePKDB(id: Long = 0L)

object SimplePKDB {
  final class SimplePKTable(tag: Tag) extends Table[SimplePKDB](tag, "simple_table") {
    def id: Rep[Long] = column[Long]("id", O.AutoInc, O.PrimaryKey)
  }
}

case class CompositePKDB(id: Long = 0L, registry: String)

object CompositePKDB {
  final class CompositePKTable(tag: Tag) extends Table[CompositePKDB](tag, "composite_table") {
    def id: Rep[Long] = column[Long]("id", O.AutoInc)
    def registry: Rep[String] = column[String]("registry")

    def pk = primaryKey(tableName + "_pk", (id, registry))
  }
}

I wish to implement a generalized distinctPK method that for querySimple: Query[SimplePKTable, SimplePKDB, Seq] unfolds to querySimple.distinctOn(_.id) and for queryComposite: Query[CompositePKTable, CompositePKDB, Seq] unfolds to queryComposite.distinctOn(row => (id, registry)).

I tried creating a trait TableWithPK and extending SimplePKTable and CompositePKTable with it

trait TableWithPK[V, PK] extends Table[V] {
  def primary: PK
}

As def primary is not Rep[_] (it is either Rep[Long], or (Rep[Long], Rep[String])), I need to provide shape: Shape[_ <: FlatShapeLevel, PK, T, _]indef distinctPK`.

def distinctPK[T <: TableWithPK[V, PK], V, PK](query: Query[T, V, Seq])(implicit shape: Shape[_ <: FlatShapeLevel, PK, T, _])
: Query[T, V, Seq] = query.distinctOn(_.primary)

However this requires implicitly carrying the shape throughout the application, which is not ideal. Is it possible to build it in the TableWithPK or extended table classes (sth like def shape: Shape[_ <: FlatShapeLevel, PK, TableWithPK[V, PK], _] inside TableWithPK)? Or this needs a completely different approach?