Skip to content

Commit e609934

Browse files
FedericoPonziFedericoPonzi
committed
A few fixes around formatting, ensure comments are printed correctly around extends/constants/assume and end of module, add more template tests
1 parent 26b916b commit e609934

26 files changed

+1027
-254
lines changed

‎src/main/java/me/fponzi/tlaplusformatter/TlaDocBuilder.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ private void registerDefaultConstructs() {
4747
registry.register(new ConstantsConstruct());
4848
registry.register(new VariableConstruct());
4949
registry.register(new OperatorConstruct());
50+
registry.register(new ConstantConstruct());
5051

5152
// Basic elements
5253
registry.register(new IdentifierConstruct());

‎src/main/java/me/fponzi/tlaplusformatter/constructs/BaseConstructFormatter.java‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ public BaseConstructFormatter(FormatConfig config) {
2929
* @param strategy Formatting strategy to use
3030
* @return Doc object for the entire construct
3131
*/
32-
protected Doc formatList(List<T> items, String prefix,
32+
protected Doc formatList(List<T> items, Doc prefix,
3333
Function<T, Doc> itemFormatter,
3434
ListFormatStrategy strategy) {
3535

@@ -38,7 +38,7 @@ protected Doc formatList(List<T> items, String prefix,
3838
}
3939

4040
if (items.size() == 1) {
41-
return Doc.text(prefix).append(itemFormatter.apply(items.get(0)));
41+
return prefix.appendSpace(itemFormatter.apply(items.get(0)));
4242
}
4343

4444
switch (strategy) {
@@ -56,8 +56,8 @@ protected Doc formatList(List<T> items, String prefix,
5656
/**
5757
* Format items as a single line with comma separation.
5858
*/
59-
protected Doc formatAsSingleLine(List<T> items, String prefix, Function<T, Doc> itemFormatter) {
60-
Doc result = Doc.text(prefix).append(itemFormatter.apply(items.get(0)));
59+
protected Doc formatAsSingleLine(List<T> items, Doc prefix, Function<T, Doc> itemFormatter) {
60+
Doc result = prefix.appendSpace(itemFormatter.apply(items.get(0)));
6161

6262
for (int i = 1; i < items.size(); i++) {
6363
result = result.append(Doc.text(", ")).append(itemFormatter.apply(items.get(i)));
@@ -70,7 +70,7 @@ protected Doc formatAsSingleLine(List<T> items, String prefix, Function<T, Doc>
7070
* Format items with smart line breaking based on line width.
7171
* Uses prettier4j's line-or-space mechanism for optimal breaking.
7272
*/
73-
protected Doc formatWithSmartBreaks(List<T> items, String prefix, Function<T, Doc> itemFormatter) {
73+
protected Doc formatWithSmartBreaks(List<T> items, Doc prefix, Function<T, Doc> itemFormatter) {
7474
Doc itemList = itemFormatter.apply(items.get(0));
7575

7676
for (int i = 1; i < items.size(); i++) {
@@ -80,21 +80,21 @@ protected Doc formatWithSmartBreaks(List<T> items, String prefix, Function<T, Do
8080
}
8181

8282
return
83-
Doc.text(prefix.trim())
84-
.appendSpace(Doc.group(itemList).indent(prefix.length()));
83+
prefix
84+
.appendSpace(Doc.group(itemList).indent(prefix.render().length()));
8585
}
8686

8787
/**
8888
* Format items with line breaks after each item.
8989
* Each item goes on its own line with proper indentation.
9090
*/
91-
protected Doc formatWithLineBreaks(List<T> items, String prefix, Function<T, Doc> itemFormatter) {
92-
Doc result = Doc.text(prefix).append(itemFormatter.apply(items.get(0)));
91+
protected Doc formatWithLineBreaks(List<T> items, Doc prefix, Function<T, Doc> itemFormatter) {
92+
Doc result = prefix.appendSpace(itemFormatter.apply(items.get(0)));
9393

9494
for (int i = 1; i < items.size(); i++) {
9595
result = result
9696
.append(Doc.text(","))
97-
.appendLine(itemFormatter.apply(items.get(i)).indent(prefix.length()));
97+
.appendLine(itemFormatter.apply(items.get(i)).indent(prefix.render().length()));
9898
}
9999

100100
return result;

‎src/main/java/me/fponzi/tlaplusformatter/constructs/ConstructContext.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,6 @@ public Doc buildChild(TreeNode child) {
4848
var comments = Arrays.stream(child.getPreComments())
4949
.map((v) -> Doc.text(v.trim()))
5050
.collect(Collectors.toList());
51-
if (!comments.isEmpty()) {
52-
System.out.println();
53-
}
5451
comments.add(childDoc);
5552
return Doc.intersperse(Doc.line(), comments);
5653
}

‎src/main/java/me/fponzi/tlaplusformatter/constructs/NodeKind.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,12 @@ public enum NodeKind {
1818
EXTENDS(350, "Extends declaration"),
1919
ACT_DECL(328, "Action declaration"),
2020
ASSUME_DECL(330, "Assume declaration"),
21-
CONS_DECL(392, "Constants declaration"),
2221
CONSTANTS(342, "Constants keyword"),
2322
VARIABLE_DECLARATION(426, "Variable declaration"),
2423
OPERATOR_DEFINITION(389, "Operator definition"),
2524
THEOREM(422, "Theorem declaration"),
2625
PARAM_DECL(391, "Parameter declaration"),
27-
PARAM_DECLARATION(392, "Parameter declaration"),
26+
PARAM_DECLARATION(392, "Parameter(==Constants) declaration"),
2827
IDENT_DECL(363, "Identifier declaration"),
2928
INFIX_DECL(370, "Infix declaration"),
3029
POSTFIX_DECL(394, "Postfix declaration"),

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/AssumptionConstruct.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,16 @@
66
import me.fponzi.tlaplusformatter.constructs.TlaConstruct;
77
import tla2sany.st.TreeNode;
88

9+
/**
10+
* Construct implementation for ASSUME/ASSUMPTION constructs.
11+
* Formats "ASSUME expr" or "ASSUMPTION expr" with proper indentation.
12+
* The node structure is expected to have two children:
13+
* - The first child is the ASSUME keyword.
14+
* - the second child is either an expression or a sequence of expressions.
15+
* <p>
16+
* Example: ASSUME x > 0
17+
* ASSUME X == <expr> (this are 4 nodes in one())
18+
*/
919
public class AssumptionConstruct implements TlaConstruct {
1020
@Override
1121
public String getName() {
@@ -22,9 +32,14 @@ public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
2232
var o = node.one();
2333
assert (o != null && o.length >= 2);
2434
var assume = context.buildChild(o[0]);
25-
return Doc.group(
26-
assume
27-
.appendSpace(context.buildChild(o[1]).indent(o[0].getImage().length() + 1))
28-
);
35+
var ret = assume.appendSpace(Doc.group(context.buildChild(o[1])).indent(o[0].getImage().length() + 1));
36+
if (o.length == 2) {
37+
return ret;
38+
}
39+
// More than one expression, need to handle line breaks
40+
// this is the case when ASSUME X == <expr>.
41+
ret = ret.appendSpace(context.buildChild(o[2])); // ==
42+
var content = Doc.group(context.buildChild(o[3])).indent(indentSize); // <expr>
43+
return ret.appendSpace(content);
2944
}
3045
}

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/ConstantConstruct.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package me.fponzi.tlaplusformatter.constructs.impl;
2+
3+
import com.opencastsoftware.prettier4j.Doc;
4+
import me.fponzi.tlaplusformatter.constructs.ConstructContext;
5+
import me.fponzi.tlaplusformatter.constructs.NodeKind;
6+
import me.fponzi.tlaplusformatter.constructs.TlaConstruct;
7+
import tla2sany.st.TreeNode;
8+
9+
/**
10+
* This is the actual "CONSTANT" or "CONSTANTS" keyword in the spec.
11+
* it has a single child which is the actual CONSTANT keyword along with comments.
12+
*/
13+
public class ConstantConstruct implements TlaConstruct {
14+
@Override
15+
public int getSupportedNodeKind() {
16+
return NodeKind.CONSTANTS.getId();
17+
}
18+
19+
@Override
20+
public String getName() {
21+
return "N_ConsDecl";
22+
}
23+
24+
@Override
25+
public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
26+
var z = node.zero();
27+
assert (z != null && z.length == 1);
28+
return context.buildChild(z[0]);
29+
}
30+
}

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/ConstantsConstruct.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,18 +18,19 @@ public class ConstantsConstruct implements TlaConstruct {
1818

1919
@Override
2020
public String getName() {
21-
return "CONSTANTS";
21+
return "N_ParamDeclaration";
2222
}
2323

2424
@Override
2525
public int getSupportedNodeKind() {
26-
return NodeKind.CONS_DECL.getId();
26+
return NodeKind.PARAM_DECLARATION.getId();
2727
}
2828

2929
@Override
3030
public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
3131
List<String> constants = context.extractStringList(node);
32-
return new ConstantsFormatter(context.getConfig()).format(constants);
32+
Doc prefix = context.buildChild(node.zero()[0]); // "CONSTANT" or "CONSTANTS" keyword
33+
return new ConstantsFormatter(context.getConfig()).format(prefix, constants);
3334
}
3435

3536
/**
@@ -41,11 +42,10 @@ public ConstantsFormatter(me.fponzi.tlaplusformatter.FormatConfig config) {
4142
super(config);
4243
}
4344

44-
public Doc format(List<String> constants) {
45+
public Doc format(Doc prefix, List<String> constants) {
4546
if (constants.isEmpty()) {
4647
return Doc.empty();
4748
}
48-
String prefix = constants.size() == 1 ? "CONSTANT " : "CONSTANTS ";
4949
ListFormatStrategy strategy = determineStrategy("CONSTANTS", constants.size());
5050
return formatList(constants, prefix, stringFormatter(), strategy);
5151
}

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/ExtendsConstruct.java‎

Lines changed: 7 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ public int getSupportedNodeKind() {
2525
@Override
2626
public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
2727
List<String> modules = context.extractStringList(node);
28-
return new ExtendsFormatter(context.getConfig()).format(modules);
28+
Doc prefix = context.buildChild(node.zero()[0]); // "EXTENDS" keyword
29+
return new ExtendsFormatter(context.getConfig()).format(prefix, modules);
2930
}
3031

3132
/**
@@ -37,16 +38,13 @@ public ExtendsFormatter(me.fponzi.tlaplusformatter.FormatConfig config) {
3738
super(config);
3839
}
3940

40-
/**
41-
* Main formatting method for EXTENDS declarations.
42-
*/
43-
public Doc format(List<String> modules) {
44-
if (modules.isEmpty()) {
41+
42+
public Doc format(Doc prefix, List<String> extendedModules) {
43+
if (extendedModules.isEmpty()) {
4544
return Doc.empty();
4645
}
47-
48-
// Use custom grouping for EXTENDS to match expected behavior
49-
return formatExtendsWithGrouping(modules);
46+
ListFormatStrategy strategy = determineStrategy("EXTENDS", extendedModules.size());
47+
return formatList(extendedModules, prefix, stringFormatter(), strategy);
5048
}
5149

5250
@Override
@@ -59,45 +57,5 @@ protected ListFormatStrategy determineStrategy(String constructName, int itemCou
5957
constructName, "breakStrategy", ListFormatStrategy.SMART_BREAK);
6058
}
6159
}
62-
63-
private String getName() {
64-
return "EXTENDS";
65-
}
66-
67-
/**
68-
* Custom formatting for EXTENDS that groups modules intelligently.
69-
*/
70-
private Doc formatExtendsWithGrouping(List<String> modules) {
71-
if (modules.isEmpty()) {
72-
return Doc.empty();
73-
}
74-
75-
if (modules.size() == 1) {
76-
return Doc.text("EXTENDS ").append(Doc.text(modules.get(0)));
77-
}
78-
79-
// For short lists, keep simple formatting
80-
if (modules.size() <= 3) {
81-
Doc moduleList = Doc.text(modules.get(0));
82-
for (int i = 1; i < modules.size(); i++) {
83-
moduleList = moduleList.append(Doc.text(", ")).append(Doc.text(modules.get(i)));
84-
}
85-
return Doc.text("EXTENDS ").append(moduleList);
86-
}
87-
88-
// For longer lists, use prettier4j line breaking
89-
Doc moduleList = Doc.text(modules.get(0));
90-
for (int i = 1; i < modules.size(); i++) {
91-
moduleList = moduleList
92-
.append(Doc.text(","))
93-
.appendLineOrSpace(Doc.text(modules.get(i)));
94-
}
95-
96-
// Use group to enable line breaking with proper indentation
97-
return Doc.group(
98-
Doc.text("EXTENDS ")
99-
.append(moduleList.indent("EXTENDS ".length()))
100-
);
101-
}
10260
}
10361
}

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/ModuleConstruct.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ public int getSupportedNodeKind() {
108108

109109
@Override
110110
public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
111-
return Doc.text(node.getHumanReadableImage());
111+
return context.buildChild(node.zero()[0]);
112112
}
113113
}
114114

‎src/main/java/me/fponzi/tlaplusformatter/constructs/impl/OpArgsConstruct.java‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,19 @@ public int getSupportedNodeKind() {
2525
public Doc buildDoc(TreeNode node, ConstructContext context, int indentSize) {
2626
var z = node.zero();
2727
assert (z != null);
28-
var ret = context.buildChild(z[0]); // (
29-
for (int i = 1; i < z.length - 1; i++) {
28+
var expr = context.buildChild(z[1]);
29+
for (int i = 2; i < z.length - 1; i++) {
3030
var doc = context.buildChild(z[i]);
31-
if (z[i].getImage().equals(",") || z[i].getImage().equals(":") || i == 1) {
32-
ret = ret.append(doc);
31+
if (z[i].getImage().equals(",") || z[i].getImage().equals(":")) {
32+
expr = expr.append(doc);
3333
} else {
34-
ret = ret.appendLineOrSpace(doc);
34+
expr = expr.appendLineOrSpace(doc);
3535
}
3636
}
37-
return ret.appendLineOrEmpty(context.buildChild(z[z.length - 1])); // )
37+
var lbracket = context.buildChild(z[0]); // (
38+
var rbracket = context.buildChild(z[z.length - 1]); // )
39+
return lbracket
40+
.append(Doc.group(expr).indent(indentSize))
41+
.appendLineOrEmpty(rbracket);
3842
}
3943
}

0 commit comments

Comments
 (0)