Attach alignment policy to each alignment group.

This supports extracting different of types alignment groups
from a single token partition range.

Alignment handler map value-type is now a generic function, allowing greater
flexibility in how alignment groups are created, replacing previous struct
entry.

Legacy adapter updated to apply one AlignmentPolicy uniformly to all alignment
groups.  This is not suitable for when we want each group to have its own
policy.

No functional change.

PiperOrigin-RevId: 331873028
diff --git a/common/formatting/align.cc b/common/formatting/align.cc
index 816bcb5..032a7e1 100644
--- a/common/formatting/align.cc
+++ b/common/formatting/align.cc
@@ -530,7 +530,7 @@
   }
 };
 
-static GroupAlignmentData AlignFilteredRows(
+static GroupAlignmentData CalculateAlignmentSpacings(
     const std::vector<TokenPartitionIterator>& rows,
     const AlignmentCellScannerFunction& cell_scanner_gen,
     MutableFormatTokenRange::iterator ftoken_base, int column_limit) {
@@ -660,15 +660,17 @@
   VLOG(1) << "end of " << __FUNCTION__;
 }
 
+// Select subset of iterators inside a partition range that are not ignored
+// by the predicate.
 static std::vector<TokenPartitionIterator> FilterAlignablePartitions(
-    const TokenPartitionRange& group,
+    const TokenPartitionRange& range,
     const IgnoreAlignmentRowPredicate& ignore_partition_predicate) {
-  // This partition group may contain partitions that should not be
+  // This partition range may contain partitions that should not be
   // considered for column alignment purposes, so filter those out.
   std::vector<TokenPartitionIterator> qualified_partitions;
-  qualified_partitions.reserve(group.size());
+  qualified_partitions.reserve(range.size());
   // like std::copy_if, but we want the iterators, not their pointees.
-  for (auto iter = group.begin(); iter != group.end(); ++iter) {
+  for (auto iter = range.begin(); iter != range.end(); ++iter) {
     if (!ignore_partition_predicate(*iter)) {
       VLOG(2) << "including partition: " << *iter;
       qualified_partitions.push_back(iter);
@@ -683,17 +685,19 @@
     const std::function<std::vector<TokenPartitionRange>(
         const TokenPartitionRange&)>& legacy_extractor,
     const IgnoreAlignmentRowPredicate& legacy_ignore_predicate,
-    const AlignmentCellScannerFunction& alignment_cell_scanner) {
-  return [legacy_extractor, legacy_ignore_predicate,
-          alignment_cell_scanner](const TokenPartitionRange& full_range) {
+    const AlignmentCellScannerFunction& alignment_cell_scanner,
+    AlignmentPolicy alignment_policy) {
+  return [legacy_extractor, legacy_ignore_predicate, alignment_cell_scanner,
+          alignment_policy](const TokenPartitionRange& full_range) {
     // must copy the closures, not just reference, to ensure valid lifetime
     const std::vector<TokenPartitionRange> ranges(legacy_extractor(full_range));
     std::vector<AlignablePartitionGroup> groups;
     groups.reserve(ranges.size());
     for (const auto& range : ranges) {
+      // Apply the same policy to all alignment groups.
       groups.emplace_back(AlignablePartitionGroup{
           FilterAlignablePartitions(range, legacy_ignore_predicate),
-          alignment_cell_scanner});
+          alignment_cell_scanner, alignment_policy});
       if (groups.back().alignable_rows.empty()) groups.pop_back();
     }
     return groups;
@@ -815,12 +819,56 @@
   return AlignmentPolicy::kPreserve;
 }
 
+static void AlignGroupUsingPolicy(
+    const AlignablePartitionGroup& alignment_group,
+    std::vector<PreFormatToken>* ftokens, absl::string_view full_text,
+    int column_limit) {
+  const TokenPartitionRange partition_range(alignment_group.Range());
+  // Compute dry-run of alignment spacings if it is needed.
+  AlignmentPolicy policy = alignment_group.alignment_policy;
+  VLOG(2) << "AlignmentPolicy: " << policy;
+  GroupAlignmentData align_data;
+  switch (policy) {
+    case AlignmentPolicy::kAlign:
+    case AlignmentPolicy::kInferUserIntent:
+      align_data =
+          CalculateAlignmentSpacings(alignment_group.alignable_rows,
+                                     alignment_group.alignment_cell_scanner,
+                                     ftokens->begin(), column_limit);
+      break;
+    default:
+      break;
+  }
+
+  // If enabled, try to decide automatically based on heurstics.
+  if (policy == AlignmentPolicy::kInferUserIntent) {
+    policy = InferUserIntendedAlignmentPolicy(partition_range, align_data);
+    VLOG(2) << "AlignmentPolicy (automatic): " << policy;
+  }
+
+  // Align or not, depending on user-elected or inferred policy.
+  switch (policy) {
+    case AlignmentPolicy::kAlign: {
+      if (!align_data.align_actions_2D.empty()) {
+        // This modifies format tokens' spacing values.
+        ApplyGroupAlignment(align_data, alignment_group, ftokens->begin());
+      }
+      break;
+    }
+    case AlignmentPolicy::kFlushLeft:
+      // This is already the default behavior elsewhere.  Nothing else to do.
+      break;
+    default:
+      IndentButPreserveOtherSpacing(partition_range, full_text, ftokens);
+      break;
+  }
+}
+
 void TabularAlignTokens(
     TokenPartitionTree* partition_ptr,
     const ExtractAlignmentGroupsFunction& extract_alignment_groups,
     std::vector<PreFormatToken>* ftokens, absl::string_view full_text,
-    const ByteOffsetSet& disabled_byte_ranges, AlignmentPolicy policy,
-    int column_limit) {
+    const ByteOffsetSet& disabled_byte_ranges, int column_limit) {
   VLOG(1) << __FUNCTION__;
   // Each subpartition is presumed to correspond to a list element or
   // possibly some other ignored element like comments.
@@ -835,9 +883,7 @@
   const std::vector<AlignablePartitionGroup> alignment_groups(
       extract_alignment_groups(subpartitions_range));
   for (const auto& alignment_group : alignment_groups) {
-    const TokenPartitionRange partition_range(
-        alignment_group.alignable_rows.front(),
-        alignment_group.alignable_rows.back() + 1);
+    const TokenPartitionRange partition_range(alignment_group.Range());
     if (partition_range.empty()) continue;
     if (AnyPartitionSubRangeIsDisabled(partition_range, full_text,
                                        disabled_byte_ranges)) {
@@ -856,42 +902,8 @@
       // TODO(b/159824483): attempt to detect and re-use pre-existing alignment
     }
 
-    // Compute dry-run of alignment spacings if it is needed.
-    VLOG(2) << "AlignmentPolicy: " << policy;
-    GroupAlignmentData align_data;
-    switch (policy) {
-      case AlignmentPolicy::kAlign:
-      case AlignmentPolicy::kInferUserIntent:
-        align_data = AlignFilteredRows(alignment_group.alignable_rows,
-                                       alignment_group.alignment_cell_scanner,
-                                       ftokens->begin(), column_limit);
-        break;
-      default:
-        break;
-    }
-
-    // If enabled, try to decide automatically based on heurstics.
-    if (policy == AlignmentPolicy::kInferUserIntent) {
-      policy = InferUserIntendedAlignmentPolicy(partition_range, align_data);
-      VLOG(2) << "AlignmentPolicy (automatic): " << policy;
-    }
-
-    // Align or not, depending on user-elected or inferred policy.
-    switch (policy) {
-      case AlignmentPolicy::kAlign: {
-        if (!align_data.align_actions_2D.empty()) {
-          // This modifies format tokens' spacing values.
-          ApplyGroupAlignment(align_data, alignment_group, ftokens->begin());
-        }
-        break;
-      }
-      case AlignmentPolicy::kFlushLeft:
-        // This is already the default behavior elsewhere.  Nothing else to do.
-        break;
-      default:
-        IndentButPreserveOtherSpacing(partition_range, full_text, ftokens);
-        break;
-    }
+    // Calculate alignment and possibly apply it depending on alignment policy.
+    AlignGroupUsingPolicy(alignment_group, ftokens, full_text, column_limit);
   }
   VLOG(1) << "end of " << __FUNCTION__;
 }
diff --git a/common/formatting/align.h b/common/formatting/align.h
index 6521aee..4d2e144 100644
--- a/common/formatting/align.h
+++ b/common/formatting/align.h
@@ -134,6 +134,36 @@
 using AlignmentCellScannerFunction =
     std::function<std::vector<ColumnPositionEntry>(const TokenPartitionTree&)>;
 
+// For sections of code that are deemed alignable, this enum controls
+// the formatter behavior.
+enum class AlignmentPolicy {
+  // Preserve text as-is.
+  kPreserve,
+
+  // No-align: flush text to left while obeying spacing constraints
+  kFlushLeft,
+
+  // Attempt tabular alignment.
+  kAlign,
+
+  // Infer whether user wanted flush-left or alignment, based on original
+  // spacing.
+  kInferUserIntent,
+};
+
+namespace internal {
+extern const std::initializer_list<
+    std::pair<const absl::string_view, AlignmentPolicy>>
+    kAlignmentPolicyNameMap;
+}  // namespace internal
+
+std::ostream& operator<<(std::ostream&, AlignmentPolicy);
+
+bool AbslParseFlag(absl::string_view text, AlignmentPolicy* policy,
+                   std::string* error);
+
+std::string AbslUnparseFlag(const AlignmentPolicy& policy);
+
 // This represents one unit of alignable work, which is usually a filtered
 // subset of partitions within a contiguous range of partitions.
 struct AlignablePartitionGroup {
@@ -143,6 +173,14 @@
   // This function scans each row to identify column positions and properties of
   // alignable cells (containing token ranges).
   AlignmentCellScannerFunction alignment_cell_scanner;
+
+  // Controls how this group should be aligned or flushed or preserved.
+  AlignmentPolicy alignment_policy;
+
+  TokenPartitionRange Range() const {
+    return TokenPartitionRange(alignable_rows.front(),
+                               alignable_rows.back() + 1);
+  }
 };
 
 // This is the interface used to sub-divide a range of token partitions into
@@ -163,7 +201,8 @@
     const std::function<std::vector<TokenPartitionRange>(
         const TokenPartitionRange&)>& legacy_extractor,
     const IgnoreAlignmentRowPredicate& legacy_ignore_predicate,
-    const AlignmentCellScannerFunction& alignment_cell_scanner);
+    const AlignmentCellScannerFunction& alignment_cell_scanner,
+    AlignmentPolicy alignment_policy);
 
 // Instantiates a ScannerType (implements ColumnSchemaScanner) and extracts
 // column alignment information, suitable as an AlignmentCellScannerFunction.
@@ -202,36 +241,6 @@
   };
 }
 
-// For sections of code that are deemed alignable, this enum controls
-// the formatter behavior.
-enum class AlignmentPolicy {
-  // Preserve text as-is.
-  kPreserve,
-
-  // No-align: flush text to left while obeying spacing constraints
-  kFlushLeft,
-
-  // Attempt tabular alignment.
-  kAlign,
-
-  // Infer whether user wanted flush-left or alignment, based on original
-  // spacing.
-  kInferUserIntent,
-};
-
-namespace internal {
-extern const std::initializer_list<
-    std::pair<const absl::string_view, AlignmentPolicy>>
-    kAlignmentPolicyNameMap;
-}  // namespace internal
-
-std::ostream& operator<<(std::ostream&, AlignmentPolicy);
-
-bool AbslParseFlag(absl::string_view text, AlignmentPolicy* policy,
-                   std::string* error);
-
-std::string AbslUnparseFlag(const AlignmentPolicy& policy);
-
 // This aligns sections of text by modifying the spacing between tokens.
 // 'partition_ptr' is a partition that can span one or more sections of
 // code to align.  The partitions themselves are not reshaped, however,
@@ -253,7 +262,6 @@
 // 'ftokens' points to the array of PreFormatTokens that spans 'full_text'.
 // 'disabled_byte_ranges' contains information about which ranges of text
 // are to preserve their original spacing (no-formatting).
-// 'policy' allows selective enabling/disabling of alignment.
 // 'column_limit' is the column width beyond which the aligner should fallback
 // to a safer action, e.g. refusing to align and leaving spacing untouched.
 //
@@ -278,8 +286,7 @@
     TokenPartitionTree* partition_ptr,
     const ExtractAlignmentGroupsFunction& extract_alignment_groups,
     std::vector<PreFormatToken>* ftokens, absl::string_view full_text,
-    const ByteOffsetSet& disabled_byte_ranges, AlignmentPolicy policy,
-    int column_limit);
+    const ByteOffsetSet& disabled_byte_ranges, int column_limit);
 
 }  // namespace verible
 
diff --git a/common/formatting/align_test.cc b/common/formatting/align_test.cc
index 501ebea..6b1dd2f 100644
--- a/common/formatting/align_test.cc
+++ b/common/formatting/align_test.cc
@@ -111,7 +111,26 @@
 static const ExtractAlignmentGroupsFunction kDefaultAlignmentHandler =
     ExtractAlignmentGroupsAdapter(
         &verible::GetSubpartitionsBetweenBlankLines, &IgnoreNone,
-        AlignmentCellScannerGenerator<TokenColumnizer>());
+        AlignmentCellScannerGenerator<TokenColumnizer>(),
+        AlignmentPolicy::kAlign);
+
+static const ExtractAlignmentGroupsFunction kFlushLeftAlignmentHandler =
+    ExtractAlignmentGroupsAdapter(
+        &verible::GetSubpartitionsBetweenBlankLines, &IgnoreNone,
+        AlignmentCellScannerGenerator<TokenColumnizer>(),
+        AlignmentPolicy::kFlushLeft);
+
+static const ExtractAlignmentGroupsFunction kPreserveAlignmentHandler =
+    ExtractAlignmentGroupsAdapter(
+        &verible::GetSubpartitionsBetweenBlankLines, &IgnoreNone,
+        AlignmentCellScannerGenerator<TokenColumnizer>(),
+        AlignmentPolicy::kPreserve);
+
+static const ExtractAlignmentGroupsFunction kInferAlignmentHandler =
+    ExtractAlignmentGroupsAdapter(
+        &verible::GetSubpartitionsBetweenBlankLines, &IgnoreNone,
+        AlignmentCellScannerGenerator<TokenColumnizer>(),
+        AlignmentPolicy::kInferUserIntent);
 
 TEST_F(TabularAlignTokenTest, EmptyPartitionRange) {
   const auto begin = pre_format_tokens_.begin();
@@ -120,7 +139,7 @@
   using tree_type = TokenPartitionTree;
   tree_type partition{all};  // no children subpartitions
   TabularAlignTokens(&partition, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
   // Not crashing is success.
   // Ideally, we would like to verify that partition was *not* modified
   // by making a deep copy and then checking DeepEqual, however,
@@ -205,7 +224,7 @@
 
 TEST_F(Sparse3x3MatrixAlignmentTest, ZeroInterTokenPadding) {
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
 
   // Sanity check: "three" (length 5) is the long-pole of the first column:
   EXPECT_EQ(pre_format_tokens_[0].before.spaces_required, tokens_[2].length());
@@ -220,8 +239,8 @@
 }
 
 TEST_F(Sparse3x3MatrixAlignmentTest, AlignmentPolicyFlushLeft) {
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kFlushLeft, 40);
+  TabularAlignTokens(&partition_, kFlushLeftAlignmentHandler,
+                     &pre_format_tokens_, sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),  // minimum spaces in this example is 0
             "onetwo\n"
@@ -234,8 +253,8 @@
   ConnectPreFormatTokensPreservedSpaceStarts(sample_.data(),
                                              &pre_format_tokens_);
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kPreserve, 40);
+  TabularAlignTokens(&partition_, kPreserveAlignmentHandler,
+                     &pre_format_tokens_, sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),  // original spacing was 1
             "one two\n"
@@ -251,7 +270,7 @@
   }
 
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -270,7 +289,7 @@
   pre_format_tokens_[4].before.spaces_required = 0;
 
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -282,7 +301,8 @@
 static const ExtractAlignmentGroupsFunction kFlushRightAlignmentHandler =
     ExtractAlignmentGroupsAdapter(
         &verible::GetSubpartitionsBetweenBlankLines, &IgnoreNone,
-        AlignmentCellScannerGenerator<TokenColumnizerRightFlushed>());
+        AlignmentCellScannerGenerator<TokenColumnizerRightFlushed>(),
+        AlignmentPolicy::kAlign);
 
 TEST_F(Sparse3x3MatrixAlignmentTest, RightFlushed) {
   // Require 1 space between tokens.
@@ -291,8 +311,7 @@
   }
 
   TabularAlignTokens(&partition_, kFlushRightAlignmentHandler,
-                     &pre_format_tokens_, sample_, ByteOffsetSet(),
-                     AlignmentPolicy::kAlign, 40);
+                     &pre_format_tokens_, sample_, ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -312,7 +331,7 @@
   }
 
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -337,9 +356,10 @@
 
   const ExtractAlignmentGroupsFunction handler = ExtractAlignmentGroupsAdapter(
       &verible::GetSubpartitionsBetweenBlankLines, ignore_threes,
-      AlignmentCellScannerGenerator<TokenColumnizer>());
+      AlignmentCellScannerGenerator<TokenColumnizer>(),
+      AlignmentPolicy::kAlign);
   TabularAlignTokens(&partition_, handler, &pre_format_tokens_, sample_,
-                     ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),         //
@@ -355,11 +375,10 @@
     ftoken.before.spaces_required = 1;
   }
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_,
-                     // Alignment disabled over entire range.
-                     ByteOffsetSet({{0, static_cast<int>(sample_.length())}}),
-                     AlignmentPolicy::kAlign, 40);
+  TabularAlignTokens(
+      &partition_, kDefaultAlignmentHandler, &pre_format_tokens_, sample_,
+      // Alignment disabled over entire range.
+      ByteOffsetSet({{0, static_cast<int>(sample_.length())}}), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -380,11 +399,10 @@
   pre_format_tokens_[2].before.break_decision = SpacingOptions::MustWrap;
   pre_format_tokens_[4].before.break_decision = SpacingOptions::MustWrap;
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_,
-                     // Alignment disabled over entire range.
-                     ByteOffsetSet({{0, static_cast<int>(sample_.length())}}),
-                     AlignmentPolicy::kAlign, 40);
+  TabularAlignTokens(
+      &partition_, kDefaultAlignmentHandler, &pre_format_tokens_, sample_,
+      // Alignment disabled over entire range.
+      ByteOffsetSet({{0, static_cast<int>(sample_.length())}}), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -422,7 +440,7 @@
       // Alignment disabled over line 2
       ByteOffsetSet({{static_cast<int>(sample_.find_first_of("\n") + 1),
                       static_cast<int>(sample_.find("four") + 4)}}),
-      AlignmentPolicy::kAlign, 40);
+      40);
 
   EXPECT_EQ(pre_format_tokens_[1].before.break_decision,
             SpacingOptions::Preserve);
@@ -446,10 +464,10 @@
   }
 
   int midpoint = sample_.length() / 2;
-  TabularAlignTokens(
-      &partition_, kDefaultAlignmentHandler, &pre_format_tokens_, sample_,
-      // Alignment disabled over partial range.
-      ByteOffsetSet({{midpoint, midpoint + 1}}), AlignmentPolicy::kAlign, 40);
+  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
+                     sample_,
+                     // Alignment disabled over partial range.
+                     ByteOffsetSet({{midpoint, midpoint + 1}}), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -465,7 +483,7 @@
   }
 
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign,
+                     sample_, ByteOffsetSet(),
                      // Column limit chosen to be smaller than sum of columns'
                      // widths. 5 (no left padding) +4 +5 = 14, so we choose 13
                      13);
@@ -488,7 +506,7 @@
 
   TabularAlignTokens(
       &partition_, kDefaultAlignmentHandler, &pre_format_tokens_, sample_,
-      ByteOffsetSet(), AlignmentPolicy::kAlign,
+      ByteOffsetSet(),
       // Column limit chosen to be smaller than sum of columns' widths.
       // 3 (indent) +5 (no left padding) +4 +5 = 17, so we choose 16
       16);
@@ -595,7 +613,7 @@
   }
 
   TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(), AlignmentPolicy::kAlign, 40);
+                     sample_, ByteOffsetSet(), 40);
 
   // Verify string rendering of result.
   EXPECT_EQ(Render(),  //
@@ -767,9 +785,8 @@
   // flushed-left looks only different by a maximum of 2 spaces ("one" vs.
   // "three", so just align it.
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(),
-                     AlignmentPolicy::kInferUserIntent, 40);
+  TabularAlignTokens(&partition_, kInferAlignmentHandler, &pre_format_tokens_,
+                     sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),      //
             "one   two\n"  //
@@ -792,9 +809,8 @@
   // The original text contains no error greater than 2 spaces relative to
   // flush-left, therefore flush-left.
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(),
-                     AlignmentPolicy::kInferUserIntent, 40);
+  TabularAlignTokens(&partition_, kInferAlignmentHandler, &pre_format_tokens_,
+                     sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),    //
             "one two\n"  //
@@ -815,9 +831,8 @@
   // The original text contains 4 excess spaces before "four" which should
   // trigger alignment.
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(),
-                     AlignmentPolicy::kInferUserIntent, 40);
+  TabularAlignTokens(&partition_, kInferAlignmentHandler, &pre_format_tokens_,
+                     sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),         //
             "one      two\n"  //
@@ -838,9 +853,8 @@
   // The original text contains only 3 excess spaces before "four" which should
   // not trigger alignment, but fall back to preserving original spacing.
 
-  TabularAlignTokens(&partition_, kDefaultAlignmentHandler, &pre_format_tokens_,
-                     sample_, ByteOffsetSet(),
-                     AlignmentPolicy::kInferUserIntent, 40);
+  TabularAlignTokens(&partition_, kInferAlignmentHandler, &pre_format_tokens_,
+                     sample_, ByteOffsetSet(), 40);
 
   EXPECT_EQ(Render(),    //
             "one two\n"  //
diff --git a/verilog/formatting/align.cc b/verilog/formatting/align.cc
index 08d7781..a1e3556 100644
--- a/verilog/formatting/align.cc
+++ b/verilog/formatting/align.cc
@@ -23,6 +23,7 @@
 
 #include "common/formatting/align.h"
 #include "common/formatting/format_token.h"
+#include "common/formatting/token_partition_tree.h"
 #include "common/formatting/unwrapped_line.h"
 #include "common/text/concrete_syntax_leaf.h"
 #include "common/text/concrete_syntax_tree.h"
@@ -41,10 +42,10 @@
 namespace verilog {
 namespace formatter {
 
+using verible::AlignablePartitionGroup;
 using verible::AlignmentCellScannerGenerator;
 using verible::AlignmentColumnProperties;
 using verible::AlignmentGroupAction;
-using verible::AlignmentPolicy;
 using verible::ByteOffsetSet;
 using verible::ColumnSchemaScanner;
 using verible::down_cast;
@@ -778,58 +779,75 @@
   bool previous_token_was_case_colon_ = false;
 };
 
-static const ExtractAlignmentGroupsFunction kPortDeclarationAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &verible::GetSubpartitionsBetweenBlankLines,
-        &IgnoreWithinPortDeclarationPartitionGroup,
-        AlignmentCellScannerGenerator<PortDeclarationColumnSchemaScanner>());
+using AlignSyntaxGroupsFunction =
+    std::function<std::vector<AlignablePartitionGroup>(
+        const TokenPartitionRange& range, const FormatStyle& style)>;
 
-static const ExtractAlignmentGroupsFunction kActualNamedParameterAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &verible::GetSubpartitionsBetweenBlankLines,
-        &IgnoreWithinActualNamedParameterPartitionGroup,
-        AlignmentCellScannerGenerator<
-            ActualNamedParameterColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignPortDeclarations(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  return verible::ExtractAlignmentGroupsAdapter(
+      &verible::GetSubpartitionsBetweenBlankLines,
+      &IgnoreWithinPortDeclarationPartitionGroup,
+      AlignmentCellScannerGenerator<PortDeclarationColumnSchemaScanner>(),
+      vstyle.port_declarations_alignment)(full_range);
+}
 
-static const ExtractAlignmentGroupsFunction kActualNamedPortAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &verible::GetSubpartitionsBetweenBlankLines,
-        &IgnoreWithinActualNamedPortPartitionGroup,
-        AlignmentCellScannerGenerator<ActualNamedPortColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignActualNamedParameters(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  return verible::ExtractAlignmentGroupsAdapter(
+      &verible::GetSubpartitionsBetweenBlankLines,
+      &IgnoreWithinActualNamedParameterPartitionGroup,
+      AlignmentCellScannerGenerator<ActualNamedParameterColumnSchemaScanner>(),
+      vstyle.named_parameter_alignment)(full_range);
+}
 
-static const ExtractAlignmentGroupsFunction kDataDeclarationAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &GetConsecutiveDataDeclarationGroups,
-        &IgnoreCommentsAndPreprocessingDirectives,
-        AlignmentCellScannerGenerator<DataDeclarationColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignActualNamedPorts(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  return verible::ExtractAlignmentGroupsAdapter(
+      &verible::GetSubpartitionsBetweenBlankLines,
+      &IgnoreWithinActualNamedPortPartitionGroup,
+      AlignmentCellScannerGenerator<ActualNamedPortColumnSchemaScanner>(),
+      vstyle.named_port_alignment)(full_range);
+}
 
-static const ExtractAlignmentGroupsFunction kClassPropertyAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &GetConsecutiveDataDeclarationGroups,
-        &IgnoreCommentsAndPreprocessingDirectives,
-        AlignmentCellScannerGenerator<ClassPropertyColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignModuleItems(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  // Currently, this only handles data/net/variable declarations.
+  // TODO(b/161814377): align continuous assignments
+  return verible::ExtractAlignmentGroupsAdapter(
+      &GetConsecutiveDataDeclarationGroups,
+      &IgnoreCommentsAndPreprocessingDirectives,
+      AlignmentCellScannerGenerator<DataDeclarationColumnSchemaScanner>(),
+      vstyle.module_net_variable_alignment)(full_range);
+}
 
-static const ExtractAlignmentGroupsFunction kCaseItemAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &verible::GetSubpartitionsBetweenBlankLines,
-        &IgnoreMultilineCaseStatements,
-        AlignmentCellScannerGenerator<CaseItemColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignClassItems(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  // TODO(fangism): align other class items besides member variables.
+  return verible::ExtractAlignmentGroupsAdapter(
+      &GetConsecutiveDataDeclarationGroups,
+      &IgnoreCommentsAndPreprocessingDirectives,
+      AlignmentCellScannerGenerator<ClassPropertyColumnSchemaScanner>(),
+      vstyle.class_member_variable_alignment)(full_range);
+}
 
-static const ExtractAlignmentGroupsFunction kParameterDeclarationAligner =
-    verible::ExtractAlignmentGroupsAdapter(
-        &verible::GetSubpartitionsBetweenBlankLines,
-        &IgnoreWithinPortDeclarationPartitionGroup,
-        AlignmentCellScannerGenerator<
-            ParameterDeclarationColumnSchemaScanner>());
+static std::vector<AlignablePartitionGroup> AlignCaseItems(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  return verible::ExtractAlignmentGroupsAdapter(
+      &verible::GetSubpartitionsBetweenBlankLines,
+      &IgnoreMultilineCaseStatements,
+      AlignmentCellScannerGenerator<CaseItemColumnSchemaScanner>(),
+      vstyle.case_items_alignment)(full_range);
+}
 
-struct AlignedFormattingConfiguration {
-  // Set of functions for driving specific code aligners.
-  ExtractAlignmentGroupsFunction handler;
-
-  // This function extracts a specific alignment policy from the
-  // Verilog-specific style structure.
-  std::function<AlignmentPolicy(const FormatStyle&)> policy;
-};
+static std::vector<AlignablePartitionGroup> AlignParameterDeclarations(
+    const TokenPartitionRange& full_range, const FormatStyle& vstyle) {
+  return verible::ExtractAlignmentGroupsAdapter(
+      &verible::GetSubpartitionsBetweenBlankLines,
+      &IgnoreWithinPortDeclarationPartitionGroup,
+      AlignmentCellScannerGenerator<ParameterDeclarationColumnSchemaScanner>(),
+      vstyle.formal_parameters_alignment)(full_range);
+}
 
 void TabularAlignTokenPartitions(TokenPartitionTree* partition_ptr,
                                  std::vector<PreFormatToken>* ftokens,
@@ -847,53 +865,28 @@
   if (node == nullptr) return;
   // Dispatch aligning function based on syntax tree node type.
 
-  static const auto* const kAlignHandlers = new std::map<
-      NodeEnum, AlignedFormattingConfiguration>{
-      {NodeEnum::kPortDeclarationList,
-       {kPortDeclarationAligner,
-        [](const FormatStyle& vstyle) {
-          return vstyle.port_declarations_alignment;
-        }}},
-      {NodeEnum::kActualParameterByNameList,
-       {kActualNamedParameterAligner,
-        [](const FormatStyle& vstyle) {
-          return vstyle.named_parameter_alignment;
-        }}},
-      {NodeEnum::kPortActualList,
-       {kActualNamedPortAligner,
-        [](const FormatStyle& vstyle) { return vstyle.named_port_alignment; }}},
-      {NodeEnum::kModuleItemList,
-       {kDataDeclarationAligner,
-        [](const FormatStyle& vstyle) {
-          return vstyle.module_net_variable_alignment;
-        }}},
-      {NodeEnum::kFormalParameterList,
-       {kParameterDeclarationAligner,
-        [](const FormatStyle& vstyle) {
-          return vstyle.formal_parameters_alignment;
-        }}},
-      {NodeEnum::kClassItems,
-       {kClassPropertyAligner,
-        [](const FormatStyle& vstyle) {
-          return vstyle.class_member_variable_alignment;
-        }}},
-      // various case-like constructs:
-      {NodeEnum::kCaseItemList,
-       {kCaseItemAligner,
-        [](const FormatStyle& vstyle) { return vstyle.case_items_alignment; }}},
-      {NodeEnum::kCaseInsideItemList,
-       {kCaseItemAligner,
-        [](const FormatStyle& vstyle) { return vstyle.case_items_alignment; }}},
-      {NodeEnum::kGenerateCaseItemList,
-       {kCaseItemAligner,
-        [](const FormatStyle& vstyle) { return vstyle.case_items_alignment; }}},
-  };
+  static const auto* const kAlignHandlers =
+      new std::map<NodeEnum, AlignSyntaxGroupsFunction>{
+          {NodeEnum::kPortDeclarationList, &AlignPortDeclarations},
+          {NodeEnum::kActualParameterByNameList, &AlignActualNamedParameters},
+          {NodeEnum::kPortActualList, &AlignActualNamedPorts},
+          {NodeEnum::kModuleItemList, &AlignModuleItems},
+          {NodeEnum::kFormalParameterList, &AlignParameterDeclarations},
+          {NodeEnum::kClassItems, &AlignClassItems},
+          // various case-like constructs:
+          {NodeEnum::kCaseItemList, &AlignCaseItems},
+          {NodeEnum::kCaseInsideItemList, &AlignCaseItems},
+          {NodeEnum::kGenerateCaseItemList, &AlignCaseItems},
+      };
   const auto handler_iter = kAlignHandlers->find(NodeEnum(node->Tag().tag));
   if (handler_iter == kAlignHandlers->end()) return;
-  verible::TabularAlignTokens(partition_ptr, handler_iter->second.handler,
-                              ftokens, full_text, disabled_byte_ranges,
-                              handler_iter->second.policy(style),
-                              style.column_limit);
+  verible::TabularAlignTokens(
+      partition_ptr,
+      [&style, handler_iter](const TokenPartitionRange& range) {
+        // Return a closure bound to the current formatting style.
+        return handler_iter->second(range, style);
+      },
+      ftokens, full_text, disabled_byte_ranges, style.column_limit);
   VLOG(1) << "end of " << __FUNCTION__;
 }