Kotlin type reference spacing is stored inconsistently and difficult to modify in recipes

[timtebeek]

Problem

Modifying the spacing around colons in Kotlin type references (e.g., val x : String, fun foo() : Int) is inconsistent and difficult for recipe authors.

Inconsistent Storage Locations

The space before the colon is stored in different places depending on context:

Context	Space stored in
Method return types	TypeReferencePrefix.prefix marker on J.MethodDeclaration
Variable declarations	JRightPadded.after of variables + TypeReferencePrefix.prefix marker
Function type parameters	TypeReferencePrefix.prefix marker on the parameter
Type parameter bounds	JContainer.before (marker prefix is EMPTY!)

Awkward Recipe API

To modify spacing around colons, recipes must:

1. Override visitMarker() and check for TypeReferencePrefix
2. Use ListUtils.map to iterate through markers
3. Know which location stores the space for each specific context
4. Handle both the marker AND other LST element spaces depending on context

There's no dedicated visitor method like visitTypeReference() that would make this discoverable and consistent.

Example: Normalizing Spacing

To normalize val x : String to val x : String, a recipe needs:

@Override
public J.VariableDeclarations visitVariableDeclarations(J.VariableDeclarations vd, ExecutionContext ctx) {
    J.VariableDeclarations v = super.visitVariableDeclarations(vd, ctx);

    // Modify the marker if present
    Markers modifiedMarkers = v.getMarkers().withMarkers(
        ListUtils.map(v.getMarkers().getMarkers(), marker -> {
            if (marker instanceof TypeReferencePrefix) {
                return ((TypeReferencePrefix) marker).withPrefix(Space.SINGLE_SPACE);
            }
            return marker;
        })
    );
    v = v.withMarkers(modifiedMarkers);

    // Also need to modify the space after the variable name
    v = v.getPadding().withVariables(
        ListUtils.map(v.getPadding().getVariables(), rpv ->
            rpv.withAfter(Space.SINGLE_SPACE)));

    // And the space after colon
    if (v.getTypeExpression() != null) {
        v = v.withTypeExpression(v.getTypeExpression().withPrefix(Space.SINGLE_SPACE));
    }
    return v;
}

This is significantly more complex than it should be.

Affected Contexts

The TypeReferencePrefix marker is used in:

• Function type parameters with names: (name: Type)
• Method return types: fun foo(): Type
• Variable declarations: val x: Type
• Type parameters with bounds: <T: Bound>
• Type constraints: where T: Comparable<T>

Proposed Solution

Add a K.TypeReference LST element that would:

• Encapsulate the prefix space (before the colon)
• Contain the type tree (with space after colon in its prefix)
• Provide a dedicated visitTypeReference() method

This would require changes to the parser, printer, visitor, and RPC serialization.

References

• Original issue: Improve quality of LST: replace applicable markers with LST elements rewrite-kotlin#477
• Test coverage added in TypeReferencePrefixTest.java

[timtebeek]

Test Coverage

The following test class documents the current behavior and demonstrates how to modify type reference spacing:

/*
 * Copyright 2026 the original author or authors.
 * <p>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * https://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openrewrite.kotlin.marker;

import org.junit.jupiter.api.Test;
import org.openrewrite.ExecutionContext;
import org.openrewrite.Issue;
import org.openrewrite.internal.ListUtils;
import org.openrewrite.java.tree.J;
import org.openrewrite.java.tree.Space;
import org.openrewrite.kotlin.KotlinIsoVisitor;
import org.openrewrite.kotlin.tree.K;
import org.openrewrite.marker.Marker;
import org.openrewrite.marker.Markers;
import org.openrewrite.test.RewriteTest;

import static org.openrewrite.kotlin.Assertions.kotlin;
import static org.openrewrite.test.RewriteTest.toRecipe;

/**
 * Tests for {@link TypeReferencePrefix} marker.
 * <p>
 * This marker is used to store the space before the colon in Kotlin type references.
 * The marker is used in several contexts:
 * <ul>
 *     <li>Function type parameters with names: {@code (name: Type)}</li>
 *     <li>Method return types: {@code fun foo(): Type}</li>
 *     <li>Variable declarations: {@code val x: Type}</li>
 *     <li>Type parameters with bounds: {@code <T: Bound>}</li>
 *     <li>Type constraints: {@code where T: Comparable<T>}</li>
 * </ul>
 * <p>
 * Note: The {@link TypeReferencePrefix} marker stores only the space BEFORE the colon.
 * The space AFTER the colon is stored in the prefix of the type element itself.
 * <p>
 * See <a href="https://github.com/openrewrite/rewrite-kotlin/issues/477">Issue 477</a>
 * for the proposal to replace this marker with a proper LST element.
 */
@Issue("https://github.com/openrewrite/rewrite-kotlin/issues/477")
@SuppressWarnings({"unused", "RedundantSuspendModifier"})
class TypeReferencePrefixTest implements RewriteTest {

    @Test
    void functionTypeParameterWithNamePreservesSpace() {
        rewriteRun(
          kotlin(
            """
              val f: (  name   :    Int ) -> Boolean = { true }
              """
          )
        );
    }

    @Test
    void methodReturnTypePreservesSpace() {
        rewriteRun(
          kotlin(
            """
              fun method ( )   :    Boolean {
                  return true
              }
              """
          )
        );
    }

    @Test
    void variableDeclarationPreservesSpace() {
        rewriteRun(
          kotlin(
            """
              val x   :    String = "hello"
              """
          )
        );
    }

    @Test
    void typeParameterBoundsPreservesSpace() {
        rewriteRun(
          kotlin(
            """
              fun < T   :    Number > method ( type : T ) { }
              """
          )
        );
    }

    @Test
    void typeConstraintPreservesSpace() {
        rewriteRun(
          kotlin(
            """
              fun <T> foo() : Int where T   :    List<T> {
                  return 0
              }
              """
          )
        );
    }

    /**
     * This test demonstrates that modifying the space before the colon in method return types
     * requires working with the {@link TypeReferencePrefix} marker directly. This is less
     * ergonomic than having a proper LST element with a dedicated visitor method.
     * <p>
     * The TypeReferencePrefix stores the space BEFORE the colon.
     * The space AFTER the colon is in the return type expression's prefix.
     */
    @Test
    void modifySpaceBeforeColonInMethodReturnType() {
        rewriteRun(
          spec -> spec.recipe(toRecipe(() -> new KotlinIsoVisitor<>() {
              @Override
              public <M extends Marker> M visitMarker(Marker marker, ExecutionContext ctx) {
                  if (marker instanceof TypeReferencePrefix) {
                      TypeReferencePrefix prefix = (TypeReferencePrefix) marker;
                      // Normalize to single space before colon
                      return (M) prefix.withPrefix(Space.SINGLE_SPACE);
                  }
                  return super.visitMarker(marker, ctx);
              }

              @Override
              public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
                  J.MethodDeclaration m = super.visitMethodDeclaration(method, ctx);
                  // Also normalize the space after colon (in the return type's prefix)
                  if (m.getReturnTypeExpression() != null) {
                      m = m.withReturnTypeExpression(m.getReturnTypeExpression().withPrefix(Space.SINGLE_SPACE));
                  }
                  return m;
              }
          })),
          kotlin(
            """
              fun method ( )      :   Boolean {
                  return true
              }
              """,
            """
              fun method ( ) : Boolean {
                  return true
              }
              """
          )
        );
    }

    /**
     * This test demonstrates that modifying the space before the colon in variable
     * declarations also requires working with the marker.
     * <p>
     * Note: For variable declarations, the space before the colon may be stored in
     * the variable's "after" space (from JRightPadded) rather than TypeReferencePrefix.
     */
    @Test
    void modifySpaceBeforeColonInVariableDeclaration() {
        rewriteRun(
          spec -> spec.recipe(toRecipe(() -> new KotlinIsoVisitor<>() {
              @Override
              public J.VariableDeclarations visitVariableDeclarations(J.VariableDeclarations vd, ExecutionContext ctx) {
                  J.VariableDeclarations v = super.visitVariableDeclarations(vd, ctx);

                  // Modify the marker if present
                  Markers modifiedMarkers = v.getMarkers().withMarkers(
                      ListUtils.map(v.getMarkers().getMarkers(), marker -> {
                          if (marker instanceof TypeReferencePrefix) {
                              return ((TypeReferencePrefix) marker).withPrefix(Space.SINGLE_SPACE);
                          }
                          return marker;
                      })
                  );
                  v = v.withMarkers(modifiedMarkers);

                  // Also need to modify the space after the variable name (before colon)
                  // This is stored in JRightPadded.after
                  v = v.getPadding().withVariables(
                      ListUtils.map(v.getPadding().getVariables(), rpv ->
                          rpv.withAfter(Space.SINGLE_SPACE)));

                  // Normalize the space after colon (in the type expression's prefix)
                  if (v.getTypeExpression() != null) {
                      v = v.withTypeExpression(v.getTypeExpression().withPrefix(Space.SINGLE_SPACE));
                  }
                  return v;
              }
          })),
          kotlin(
            """
              val x      :   String = "hello"
              """,
            """
              val x : String = "hello"
              """
          )
        );
    }

    /**
     * This test demonstrates modifying the space in function type parameters.
     * <p>
     * For function type parameters:
     * - TypeReferencePrefix stores space before the colon (after the name)
     * - The parameterType's prefix stores the space after the colon
     */
    @Test
    void modifySpaceBeforeColonInFunctionTypeParameter() {
        rewriteRun(
          spec -> spec.recipe(toRecipe(() -> new KotlinIsoVisitor<>() {
              @Override
              public K.FunctionType.Parameter visitFunctionTypeParameter(K.FunctionType.Parameter parameter, ExecutionContext ctx) {
                  K.FunctionType.Parameter p = super.visitFunctionTypeParameter(parameter, ctx);
                  // Modify the marker on the parameter using withMarkers and ListUtils.map
                  Markers modifiedMarkers = p.getMarkers().withMarkers(
                      ListUtils.map(p.getMarkers().getMarkers(), marker -> {
                          if (marker instanceof TypeReferencePrefix) {
                              return ((TypeReferencePrefix) marker).withPrefix(Space.SINGLE_SPACE);
                          }
                          return marker;
                      })
                  );
                  p = p.withMarkers(modifiedMarkers);
                  // Also modify the space after the colon (in parameterType's prefix)
                  p = p.withParameterType(p.getParameterType().withPrefix(Space.SINGLE_SPACE));
                  return p;
              }
          })),
          kotlin(
            """
              val f: (  name      :    Int ) -> Boolean = { true }
              """,
            """
              val f: (  name : Int ) -> Boolean = { true }
              """
          )
        );
    }

    /**
     * Test that type parameter bounds preserve space correctly.
     * <p>
     * For type parameters with bounds, the spacing is stored in multiple places:
     * - TypeReferencePrefix is a marker indicating a colon should be printed (often with empty space)
     * - The bounds container's "before" space stores the space between name and colon
     * - The first bound's prefix stores the space after the colon
     * <p>
     * This is an example of how the TypeReferencePrefix marker makes formatting
     * more complex than necessary - the space is split across multiple locations.
     */
    @Test
    void modifySpaceBeforeColonInTypeParameterBounds() {
        rewriteRun(
          spec -> spec.recipe(toRecipe(() -> new KotlinIsoVisitor<>() {
              @Override
              public J.TypeParameter visitTypeParameter(J.TypeParameter typeParam, ExecutionContext ctx) {
                  J.TypeParameter tp = super.visitTypeParameter(typeParam, ctx);

                  // For type parameters, the TypeReferencePrefix.prefix is typically EMPTY
                  // The actual space before the colon is in bounds container's "before"

                  // Modify the space before the colon (in bounds container's "before")
                  if (tp.getPadding().getBounds() != null) {
                      tp = tp.getPadding().withBounds(
                          tp.getPadding().getBounds().withBefore(Space.SINGLE_SPACE));
                  }

                  // Also modify the space after the colon (in first bound's prefix)
                  if (tp.getBounds() != null && !tp.getBounds().isEmpty()) {
                      tp = tp.withBounds(ListUtils.map(tp.getBounds(), b ->
                          b.withPrefix(Space.SINGLE_SPACE)));
                  }
                  return tp;
              }
          })),
          kotlin(
            """
              fun < T      :   Number > method ( type : T ) { }
              """,
            // Note: The colon is attached to the type name without space before it
            // because the TypeReferencePrefix marker has Space.EMPTY and we can only
            // modify bounds.before. The colon is printed AFTER bounds.before.
            """
              fun < T : Number > method ( type : T ) { }
              """
          )
        );
    }
}