如何在Chisel中将状态机拆分成多个类或特征？

import chisel3._

class StateMachine extends Module {
  val io = IO(new Bundle {
    val input = Input(UInt(1.W))
    val output = Output(UInt(1.W))
  })
  
  // 主类中的顶层逻辑和控制流程
  val state = RegInit(0.U(2.W))
  
  // 实例化拆分出来的类或特征
  val stateA = Module(new StateA)
  val stateB = Module(new StateB)
  
  // 连接拆分出来的类或特征
  stateA.io.input := io.input
  stateA.io.state := state
  stateB.io.state := state
  
  // 主类根据需要调用拆分出来的类或特征的方法
  when(state === 0.U) {
    stateA.io.enable := true.B
    stateB.io.enable := false.B
    state := stateA.io.nextState
  }.elsewhen(state === 1.U) {
    stateA.io.enable := false.B
    stateB.io.enable := true.B
    state := stateB.io.nextState
  }.otherwise {
    stateA.io.enable := false.B
    stateB.io.enable := false.B
    state := 0.U
  }
  
  // 主类连接拆分出来的类或特征的输出
  io.output := stateA.io.output | stateB.io.output
}

// 拆分出来的类或特征之一
class StateA extends Module {
  val io = IO(new Bundle {
    val input = Input(UInt(1.W))
    val state = Input(UInt(2.W))
    val enable = Input(Bool())
    val nextState = Output(UInt(2.W))
    val output = Output(UInt(1.W))
  })
  
  // 类或特征的逻辑和控制流程
  val nextState = RegNext(io.state)
  val output = io.input
  
  when(io.enable) {
    nextState := io.state + 1.U
  }
}

// 拆分出来的类或特征之二
class StateB extends Module {
  val io = IO(new Bundle {
    val state = Input(UInt(2.W))
    val enable = Input(Bool())
    val nextState = Output(UInt(2.W))
    val output = Output(UInt(1.W))
  })
  
  // 类或特征的逻辑和控制流程
  val nextState = RegNext(io.state)
  val output = io.state(0)
  
  when(io.enable) {
    nextState := io.state - 1.U
  }
}