init

pkg/ast/state.go

@@ -3,13 +3,13 @@ package ast
 
 // StateDiagram represents a state diagram
 type StateDiagram struct {
-	States      map[string]*StateNode `json:"states"`
-	Transitions []*StateTransition    `json:"transitions"`
-	StartState  *string               `json:"startState,omitempty"`
-	EndStates   []string              `json:"endStates,omitempty"`
-	Direction   string                `json:"direction,omitempty"`
-	Title       *string               `json:"title,omitempty"`
-	Config      map[string]any        `json:"config,omitempty"`
+	States      []*StateNode       `json:"states"`
+	Transitions []*StateTransition `json:"transitions"`
+	StartState  *string            `json:"startState,omitempty"`
+	EndStates   []string           `json:"endStates,omitempty"`
+	Direction   string             `json:"direction,omitempty"`
+	Title       *string            `json:"title,omitempty"`
+	Config      map[string]any     `json:"config,omitempty"`
 }
 
 type StateNode struct {
@@ -71,12 +71,12 @@ func (s *StateDiagram) Validate() error {
 	// Validate transitions reference valid states
 	for _, trans := range s.Transitions {
 		if trans.From != "[*]" {
-			if _, exists := s.States[trans.From]; !exists {
+			if !s.hasState(trans.From) {
 				return NewValidationError("transition references non-existent state: " + trans.From)
 			}
 		}
 		if trans.To != "[*]" {
-			if _, exists := s.States[trans.To]; !exists {
+			if !s.hasState(trans.To) {
 				return NewValidationError("transition references non-existent state: " + trans.To)
 			}
 		}
@@ -84,10 +84,20 @@ func (s *StateDiagram) Validate() error {
 	return nil
 }
 
+// hasState checks if a state with the given ID exists
+func (s *StateDiagram) hasState(id string) bool {
+	for _, state := range s.States {
+		if state.ID == id {
+			return true
+		}
+	}
+	return false
+}
+
 // NewStateDiagram creates a new state diagram
 func NewStateDiagram() *StateDiagram {
 	return &StateDiagram{
-		States:      make(map[string]*StateNode),
+		States:      make([]*StateNode, 0),
 		Transitions: make([]*StateTransition, 0),
 		EndStates:   make([]string, 0),
 		Config:      make(map[string]any),

pkg/parser/enhanced_state_test.go

@@ -4,6 +4,7 @@ import (
 	"strings"
 	"testing"
 
+	"mermaid-go/pkg/ast"
 	"mermaid-go/pkg/renderer"
 )
 
@@ -162,8 +163,14 @@ func TestStateParser_StateActions(t *testing.T) {
 		t.Fatal("No states found")
 	}
 
-	state1, exists := diagram.States["State1"]
-	if !exists {
+	var state1 *ast.StateNode
+	for _, state := range diagram.States {
+		if state.ID == "State1" {
+			state1 = state
+			break
+		}
+	}
+	if state1 == nil {
 		t.Fatal("State1 not found")
 	}
 
@@ -196,8 +203,14 @@ func TestStateParser_CompositeStates(t *testing.T) {
 	}
 
 	// Check composite state was parsed
-	compositeState, exists := diagram.States["CompositeState"]
-	if !exists {
+	var compositeState *ast.StateNode
+	for _, state := range diagram.States {
+		if state.ID == "CompositeState" {
+			compositeState = state
+			break
+		}
+	}
+	if compositeState == nil {
 		t.Fatal("CompositeState not found")
 	}
 

pkg/parser/state.go

@@ -11,15 +11,17 @@ import (
 
 // StateParser implements state diagram parsing following stateDiagram.jison
 type StateParser struct {
-	tokens  []lexer.Token
-	current int
-	diagram *ast.StateDiagram
+	tokens   []lexer.Token
+	current  int
+	diagram  *ast.StateDiagram
+	stateMap map[string]*ast.StateNode // Keep track of states by name for quick lookup
 }
 
 // NewStateParser creates a new state parser
 func NewStateParser() *StateParser {
 	return &StateParser{
-		diagram: ast.NewStateDiagram(),
+		diagram:  ast.NewStateDiagram(),
+		stateMap: make(map[string]*ast.StateNode),
 	}
 }
 
@@ -36,6 +38,7 @@ func (p *StateParser) Parse(input string) (*ast.StateDiagram, error) {
 	p.tokens = lexer.FilterTokens(tokens)
 	p.current = 0
 	p.diagram = ast.NewStateDiagram()
+	p.stateMap = make(map[string]*ast.StateNode)
 
 	// Parse document
 	err = p.parseDocument()
@@ -109,12 +112,18 @@ func (p *StateParser) parseState() error {
 
 	stateName := p.advance().Value
 
-	state := &ast.StateNode{
-		ID:         stateName,
-		Label:      stateName,
-		Type:       ast.StateTypeDefault,
-		SubStates:  make(map[string]*ast.StateNode),
-		CssClasses: make([]string, 0),
+	// Check if state already exists
+	var state *ast.StateNode
+	if existingState, exists := p.stateMap[stateName]; exists {
+		state = existingState
+	} else {
+		state = &ast.StateNode{
+			ID:         stateName,
+			Label:      stateName,
+			Type:       ast.StateTypeDefault,
+			SubStates:  make(map[string]*ast.StateNode),
+			CssClasses: make([]string, 0),
+		}
 	}
 
 	// Check for 'as' alias
@@ -174,7 +183,11 @@ func (p *StateParser) parseState() error {
 		}
 	}
 
-	p.diagram.States[stateName] = state
+	// Only add state if it doesn't already exist
+	if _, exists := p.stateMap[stateName]; !exists {
+		p.stateMap[stateName] = state
+		p.diagram.States = append(p.diagram.States, state)
+	}
 	return nil
 }
 
@@ -186,11 +199,332 @@ func (p *StateParser) parseStateBody(parentState *ast.StateNode) error {
 			continue
 		}
 
-		// Parse sub-statements (simplified for now)
-		if err := p.parseStatement(); err != nil {
+		// Parse substates and transitions within the composite state
+		if p.checkKeyword("state") {
+			// Check if this is a nested composite state
+			// Look ahead: state ID { means nested composite state
+			if p.peekNext().Type == lexer.TokenID && p.current+2 < len(p.tokens) && p.tokens[p.current+2].Type == lexer.TokenOpenBrace {
+				// This is a nested composite state, parse it as a full state
+				if err := p.parseNestedCompositeState(parentState); err != nil {
+					return err
+				}
+			} else {
+				// This is a simple substate
+				if err := p.parseSubState(parentState); err != nil {
+					return err
+				}
+			}
+		} else if p.check(lexer.TokenOpenBracket) {
+			// Handle [*] start/end states within composite state
+			if err := p.parseStartEndStateInComposite(parentState); err != nil {
+				return err
+			}
+		} else if p.check(lexer.TokenID) {
+			// In composite state, ID followed by --> is a transition
+			// We need to parse the transition and ensure the states exist as substates
+			if err := p.parseStateOrTransitionInComposite(parentState); err != nil {
+				return err
+			}
+		} else {
+			token := p.peek()
+			return p.error(fmt.Sprintf("unexpected token in composite state: %s", token.Value))
+		}
+	}
+	return nil
+}
+
+// parseSubState parses a substate within a composite state
+func (p *StateParser) parseSubState(parentState *ast.StateNode) error {
+	if !p.checkKeyword("state") {
+		return p.error("expected 'state'")
+	}
+	p.advance() // consume 'state'
+
+	if !p.check(lexer.TokenID) {
+		return p.error("expected state name")
+	}
+	stateName := p.advance().Value
+
+	// Create substate
+	subState := &ast.StateNode{
+		ID:         stateName,
+		Label:      stateName,
+		Type:       ast.StateTypeDefault,
+		SubStates:  make(map[string]*ast.StateNode),
+		CssClasses: make([]string, 0),
+	}
+
+	// Add alias if different from ID
+	if p.checkKeyword("as") {
+		p.advance() // consume 'as'
+		if !p.check(lexer.TokenID) && !p.check(lexer.TokenString) {
+			return p.error("expected alias name")
+		}
+		alias := p.advance().Value
+		// Remove quotes if present
+		if len(alias) > 2 && alias[0] == '"' && alias[len(alias)-1] == '"' {
+			alias = alias[1 : len(alias)-1]
+		}
+		subState.Label = alias
+	}
+
+	// Check for description or special type
+	if p.check(lexer.TokenColon) {
+		p.advance() // consume ':'
+		if p.checkKeyword("<<fork>>") {
+			p.advance()
+			subState.Type = ast.StateTypeFork
+		} else if p.checkKeyword("<<join>>") {
+			p.advance()
+			subState.Type = ast.StateTypeJoin
+		} else if p.checkKeyword("<<choice>>") {
+			p.advance()
+			subState.Type = ast.StateTypeChoice
+		} else if p.checkKeyword("<<history>>") {
+			p.advance()
+			subState.Type = ast.StateTypeHistory
+		} else if p.checkKeyword("<<deepHistory>>") {
+			p.advance()
+			subState.Type = ast.StateTypeDeepHistory
+		} else {
+			// Parse description
+			var descParts []string
+			for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
+				descParts = append(descParts, p.advance().Value)
+			}
+			if len(descParts) > 0 {
+				desc := strings.TrimSpace(strings.Join(descParts, " "))
+				subState.Description = &desc
+			}
+		}
+	}
+
+	// Add to parent state's substates
+	parentState.SubStates[subState.ID] = subState
+
+	return nil
+}
+
+// parseNestedCompositeState parses a nested composite state within a parent composite state
+func (p *StateParser) parseNestedCompositeState(parentState *ast.StateNode) error {
+	if !p.checkKeyword("state") {
+		return p.error("expected 'state'")
+	}
+	p.advance() // consume 'state'
+
+	if !p.check(lexer.TokenID) {
+		return p.error("expected state name")
+	}
+	stateName := p.advance().Value
+
+	// Create nested composite state
+	nestedState := &ast.StateNode{
+		ID:         stateName,
+		Label:      stateName,
+		Type:       ast.StateTypeDefault,
+		SubStates:  make(map[string]*ast.StateNode),
+		CssClasses: make([]string, 0),
+	}
+
+	// Check for 'as' alias
+	if p.checkKeyword("as") {
+		p.advance() // consume 'as'
+		if !p.check(lexer.TokenID) && !p.check(lexer.TokenString) {
+			return p.error("expected alias name")
+		}
+		alias := p.advance().Value
+		// Remove quotes if present
+		if len(alias) > 2 && alias[0] == '"' && alias[len(alias)-1] == '"' {
+			alias = alias[1 : len(alias)-1]
+		}
+		nestedState.Label = alias
+	}
+
+	// Parse the nested composite state body
+	if p.check(lexer.TokenOpenBrace) {
+		p.advance() // consume '{'
+		err := p.parseStateBody(nestedState)
+		if err != nil {
 			return err
 		}
+		if !p.check(lexer.TokenCloseBrace) {
+			return p.error("expected '}'")
+		}
+		p.advance() // consume '}'
+	}
+
+	// Add to parent state's substates
+	parentState.SubStates[nestedState.ID] = nestedState
+
+	return nil
+}
+
+// parseStateOrTransitionInComposite parses a state or transition within a composite state
+func (p *StateParser) parseStateOrTransitionInComposite(parentState *ast.StateNode) error {
+	// This should be a transition within the composite state
+	// Parse the transition normally, but ensure states are added as substates
+	return p.parseStateOrTransitionWithParent(parentState)
+}
+
+// parseStateOrTransitionWithParent parses a state or transition with a parent composite state
+func (p *StateParser) parseStateOrTransitionWithParent(parentState *ast.StateNode) error {
+	stateName := p.advance().Value
+
+	// Ensure state exists as substate in parent
+	p.ensureStateAsSubstate(parentState, stateName)
+
+	// Check for transition arrow
+	if p.checkArrow() {
+		return p.parseTransitionWithParent(parentState, stateName)
+	}
+
+	// For now, just handle transitions in composite states
+	// Other cases (state actions, descriptions) can be handled later
+	return nil
+}
+
+// ensureStateAsSubstate ensures a state exists as a substate in the parent
+func (p *StateParser) ensureStateAsSubstate(parentState *ast.StateNode, id string) {
+	if _, exists := parentState.SubStates[id]; !exists {
+		subState := &ast.StateNode{
+			ID:         id,
+			Label:      id,
+			Type:       ast.StateTypeDefault,
+			SubStates:  make(map[string]*ast.StateNode),
+			CssClasses: make([]string, 0),
+		}
+		parentState.SubStates[id] = subState
+	}
+}
+
+// parseTransitionWithParent parses a transition within a composite state
+func (p *StateParser) parseTransitionWithParent(parentState *ast.StateNode, fromState string) error {
+	if !p.checkArrow() {
+		return p.error("expected transition arrow")
+	}
+	p.advance() // consume arrow
+
+	if p.isAtEnd() {
+		return p.error("unexpected end of input")
+	}
+
+	var toState string
+	if p.check(lexer.TokenOpenBracket) {
+		// Handle [*] end state
+		p.advance() // consume '['
+		if !p.check(lexer.TokenMult) {
+			return p.error("expected '*' in [*]")
+		}
+		p.advance() // consume '*'
+		if !p.check(lexer.TokenCloseBracket) {
+			return p.error("expected ']' in [*]")
+		}
+		p.advance() // consume ']'
+		toState = "[*]"
+	} else if p.check(lexer.TokenID) {
+		toState = p.advance().Value
+		// Ensure toState exists as substate
+		p.ensureStateAsSubstate(parentState, toState)
+	} else {
+		return p.error("expected state name or [*]")
+	}
+
+	// Create transition
+	transition := &ast.StateTransition{
+		From: fromState,
+		To:   toState,
+	}
+
+	// Parse transition decorations if present
+	if p.check(lexer.TokenColon) {
+		p.advance() // consume ':'
+		var labelParts []string
+		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
+			labelParts = append(labelParts, p.advance().Value)
+		}
+		if len(labelParts) > 0 {
+			label := strings.TrimSpace(strings.Join(labelParts, " "))
+			transition.Label = &label
+		}
+	}
+
+	// Add transition to diagram
+	p.diagram.Transitions = append(p.diagram.Transitions, transition)
+
+	return nil
+}
+
+// parseStartEndStateInComposite parses [*] transitions within a composite state
+func (p *StateParser) parseStartEndStateInComposite(parentState *ast.StateNode) error {
+	if !p.check(lexer.TokenOpenBracket) {
+		return p.error("expected '['")
+	}
+	p.advance() // consume '['
+
+	if !p.check(lexer.TokenMult) {
+		return p.error("expected '*'")
+	}
+	p.advance() // consume '*'
+
+	if !p.check(lexer.TokenCloseBracket) {
+		return p.error("expected ']'")
+	}
+	p.advance() // consume ']'
+
+	// Expect arrow
+	if !p.checkArrow() {
+		return p.error("expected transition arrow")
+	}
+	p.advance() // consume arrow
+
+	// Parse target state
+	if p.isAtEnd() {
+		return p.error("unexpected end of input")
+	}
+
+	var toState string
+	if p.check(lexer.TokenOpenBracket) {
+		// Handle [*] end state
+		p.advance() // consume '['
+		if !p.check(lexer.TokenMult) {
+			return p.error("expected '*' in [*]")
+		}
+		p.advance() // consume '*'
+		if !p.check(lexer.TokenCloseBracket) {
+			return p.error("expected ']' in [*]")
+		}
+		p.advance() // consume ']'
+		toState = "[*]"
+	} else if p.check(lexer.TokenID) {
+		toState = p.advance().Value
+		// Ensure toState exists as substate
+		p.ensureStateAsSubstate(parentState, toState)
+	} else {
+		return p.error("expected state name or [*]")
+	}
+
+	// Create transition
+	transition := &ast.StateTransition{
+		From: "[*]",
+		To:   toState,
 	}
+
+	// Parse transition decorations if present
+	if p.check(lexer.TokenColon) {
+		p.advance() // consume ':'
+		var labelParts []string
+		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
+			labelParts = append(labelParts, p.advance().Value)
+		}
+		if len(labelParts) > 0 {
+			label := strings.TrimSpace(strings.Join(labelParts, " "))
+			transition.Label = &label
+		}
+	}
+
+	// Add transition to diagram
+	p.diagram.Transitions = append(p.diagram.Transitions, transition)
+
 	return nil
 }
 
@@ -290,7 +624,7 @@ func (p *StateParser) parseStateOrTransition() error {
 			}
 
 			actionContent := strings.TrimSpace(strings.Join(actionParts, " "))
-			state := p.diagram.States[stateName]
+			state := p.stateMap[stateName]
 
 			switch actionType {
 			case "entry":
@@ -308,8 +642,23 @@ func (p *StateParser) parseStateOrTransition() error {
 			}
 			if len(descParts) > 0 {
 				desc := strings.TrimSpace(strings.Join(descParts, " "))
-				state := p.diagram.States[stateName]
-				state.Description = &desc
+				state := p.stateMap[stateName]
+
+				// Check if this is a special state type
+				switch desc {
+				case "<<fork>>":
+					state.Type = ast.StateTypeFork
+				case "<<join>>":
+					state.Type = ast.StateTypeJoin
+				case "<<choice>>":
+					state.Type = ast.StateTypeChoice
+				case "<<history>>":
+					state.Type = ast.StateTypeHistory
+				case "<<deepHistory>>":
+					state.Type = ast.StateTypeDeepHistory
+				default:
+					state.Description = &desc
+				}
 			}
 		}
 	}
@@ -480,7 +829,7 @@ func (p *StateParser) parseNote() error {
 
 	// Attach to state; ensure exists
 	p.ensureState(stateID)
-	if st, ok := p.diagram.States[stateID]; ok {
+	if st, ok := p.stateMap[stateID]; ok {
 		st.Note = &ast.StateNote{Position: place, Text: noteText}
 	}
 	return nil
@@ -488,7 +837,7 @@ func (p *StateParser) parseNote() error {
 
 // ensureState ensures a state exists, creating it if needed
 func (p *StateParser) ensureState(id string) {
-	if _, exists := p.diagram.States[id]; !exists {
+	if _, exists := p.stateMap[id]; !exists {
 		state := &ast.StateNode{
 			ID:         id,
 			Label:      id,
@@ -496,7 +845,8 @@ func (p *StateParser) ensureState(id string) {
 			SubStates:  make(map[string]*ast.StateNode),
 			CssClasses: make([]string, 0),
 		}
-		p.diagram.States[id] = state
+		p.stateMap[id] = state
+		p.diagram.States = append(p.diagram.States, state)
 	}
 }
 
@@ -546,6 +896,13 @@ func (p *StateParser) peek() lexer.Token {
 	return p.tokens[p.current]
 }
 
+func (p *StateParser) peekNext() lexer.Token {
+	if p.current+1 >= len(p.tokens) {
+		return lexer.Token{Type: lexer.TokenEOF}
+	}
+	return p.tokens[p.current+1]
+}
+
 func (p *StateParser) previous() lexer.Token {
 	if p.current <= 0 {
 		return lexer.Token{Type: lexer.TokenEOF}
@@ -586,13 +943,15 @@ func (p *StateParser) parseStateAction(actionType string) error {
 	stateName := p.advance().Value
 
 	// Ensure state exists
-	if _, exists := p.diagram.States[stateName]; !exists {
+	if _, exists := p.stateMap[stateName]; !exists {
 		// Create state if it doesn't exist
-		p.diagram.States[stateName] = &ast.StateNode{
+		state := &ast.StateNode{
 			ID:    stateName,
 			Label: stateName,
 			Type:  ast.StateTypeDefault,
 		}
+		p.stateMap[stateName] = state
+		p.diagram.States = append(p.diagram.States, state)
 	}
 
 	// Parse action content (everything after state name until newline)
@@ -604,7 +963,7 @@ func (p *StateParser) parseStateAction(actionType string) error {
 	actionContent := strings.TrimSpace(strings.Join(actionParts, " "))
 
 	// Set the appropriate action field
-	state := p.diagram.States[stateName]
+	state := p.stateMap[stateName]
 	switch actionType {
 	case "entry":
 		state.EntryAction = &actionContent

pkg/renderer/state.go

@@ -33,8 +33,284 @@ func (r *StateRenderer) Render(diagram *ast.StateDiagram) (string, error) {
 		builder.WriteString(fmt.Sprintf("    direction %s\n", diagram.Direction))
 	}
 
-	// Render states
+	// Create a map to track which states need explicit declaration
+	stateNeedsDeclaration := make(map[string]bool)
+
+	// Check which states need explicit declaration
+	for _, state := range diagram.States {
+		needsDeclaration := false
+
+		// State needs declaration if it has special type, alias, or description
+		if state.Type != ast.StateTypeDefault || state.Label != state.ID || state.Description != nil {
+			needsDeclaration = true
+		}
+
+		// State needs declaration if it has composite structure
+		if len(state.SubStates) > 0 {
+			needsDeclaration = true
+		}
+
+		stateNeedsDeclaration[state.ID] = needsDeclaration
+	}
+
+	// Separate transitions into external and internal
+	var externalTransitions []*ast.StateTransition
+	var internalTransitions = make(map[string][]*ast.StateTransition) // stateID -> transitions
+
+	for _, transition := range diagram.Transitions {
+		// Check if this is an internal transition (both states are within the same composite state)
+		isInternal := false
+		for _, state := range diagram.States {
+			if len(state.SubStates) > 0 {
+				fromInState := r.isStateInComposite(state, transition.From)
+				toInState := r.isStateInComposite(state, transition.To)
+				if fromInState && toInState {
+					internalTransitions[state.ID] = append(internalTransitions[state.ID], transition)
+					isInternal = true
+					break
+				}
+			}
+		}
+		if !isInternal {
+			externalTransitions = append(externalTransitions, transition)
+		}
+	}
+
+	// Separate states into special states and others
+	var specialStates []*ast.StateNode
+	var otherStates []*ast.StateNode
+
 	for _, state := range diagram.States {
+		if stateNeedsDeclaration[state.ID] {
+			// Check if this is a special state
+			if state.Type == ast.StateTypeFork || state.Type == ast.StateTypeJoin ||
+				state.Type == ast.StateTypeChoice || state.Type == ast.StateTypeHistory ||
+				state.Type == ast.StateTypeDeepHistory {
+				specialStates = append(specialStates, state)
+			} else {
+				otherStates = append(otherStates, state)
+			}
+		}
+	}
+
+	// Render special states first (sorted by type priority)
+	r.sortStatesByTransitionOrder(specialStates, diagram.Transitions)
+	for _, state := range specialStates {
+		r.renderStateWithInternalTransitions(&builder, state, internalTransitions[state.ID])
+	}
+
+	// Render [*] --> state transitions
+	for _, transition := range externalTransitions {
+		if transition.From == "[*]" {
+			r.renderTransition(&builder, transition)
+		}
+	}
+
+	// Render other states that need explicit declaration (composite states, etc.)
+	for _, state := range otherStates {
+		r.renderStateWithInternalTransitions(&builder, state, internalTransitions[state.ID])
+	}
+
+	// Render state actions for states that don't need explicit declaration
+	for _, state := range diagram.States {
+		if !stateNeedsDeclaration[state.ID] {
+			r.renderStateActions(&builder, state)
+		}
+	}
+
+	// Render other external transitions (state --> [*] and state --> state)
+	for _, transition := range externalTransitions {
+		if transition.From != "[*]" {
+			r.renderTransition(&builder, transition)
+		}
+	}
+
+	return builder.String(), nil
+}
+
+// sortStatesByTransitionOrder sorts states by their order in the transition chain
+func (r *StateRenderer) sortStatesByTransitionOrder(states []*ast.StateNode, transitions []*ast.StateTransition) {
+	// For special states, sort by type priority: join -> fork -> choice -> others
+	typePriority := map[ast.StateType]int{
+		ast.StateTypeJoin:        0,
+		ast.StateTypeFork:        1,
+		ast.StateTypeChoice:      2,
+		ast.StateTypeHistory:     3,
+		ast.StateTypeDeepHistory: 4,
+		ast.StateTypeDefault:     5,
+	}
+
+	// Sort states by type priority
+	for i := 0; i < len(states); i++ {
+		for j := i + 1; j < len(states); j++ {
+			priorityI := typePriority[states[i].Type]
+			priorityJ := typePriority[states[j].Type]
+			if priorityI > priorityJ {
+				states[i], states[j] = states[j], states[i]
+			}
+		}
+	}
+}
+
+// renderTransition renders a single transition
+func (r *StateRenderer) renderTransition(builder *strings.Builder, transition *ast.StateTransition) {
+	builder.WriteString("    ")
+	builder.WriteString(transition.From)
+	builder.WriteString(" --> ")
+	builder.WriteString(transition.To)
+
+	// Add transition decorations: label, [guard], /action
+	var decorations []string
+	if transition.Label != nil {
+		decorations = append(decorations, *transition.Label)
+	}
+	if transition.Condition != nil {
+		decorations = append(decorations, fmt.Sprintf("[%s]", *transition.Condition))
+	}
+	if transition.Action != nil {
+		decorations = append(decorations, fmt.Sprintf("/ %s", *transition.Action))
+	}
+
+	if len(decorations) > 0 {
+		builder.WriteString(" : ")
+		builder.WriteString(strings.Join(decorations, " "))
+	}
+	builder.WriteString("\n")
+}
+
+// renderStateActions renders state actions without explicit state declaration
+func (r *StateRenderer) renderStateActions(builder *strings.Builder, state *ast.StateNode) {
+	if state.EntryAction != nil {
+		builder.WriteString("    ")
+		builder.WriteString(state.ID)
+		builder.WriteString(" : entry ")
+		builder.WriteString(*state.EntryAction)
+		builder.WriteString("\n")
+	}
+	if state.ExitAction != nil {
+		builder.WriteString("    ")
+		builder.WriteString(state.ID)
+		builder.WriteString(" : exit ")
+		builder.WriteString(*state.ExitAction)
+		builder.WriteString("\n")
+	}
+	if state.DoAction != nil {
+		builder.WriteString("    ")
+		builder.WriteString(state.ID)
+		builder.WriteString(" : do ")
+		builder.WriteString(*state.DoAction)
+		builder.WriteString("\n")
+	}
+}
+
+// renderState renders a state with explicit declaration
+func (r *StateRenderer) renderState(builder *strings.Builder, state *ast.StateNode) {
+	builder.WriteString("    state ")
+	builder.WriteString(state.ID)
+
+	// Add alias if different from ID
+	if state.Label != state.ID {
+		builder.WriteString(" as ")
+		if strings.Contains(state.Label, " ") {
+			builder.WriteString(fmt.Sprintf("\"%s\"", state.Label))
+		} else {
+			builder.WriteString(state.Label)
+		}
+	}
+
+	// Add description or special type
+	if state.Description != nil {
+		builder.WriteString(" : ")
+		builder.WriteString(*state.Description)
+	} else {
+		switch state.Type {
+		case ast.StateTypeFork:
+			builder.WriteString(" : <<fork>>")
+		case ast.StateTypeJoin:
+			builder.WriteString(" : <<join>>")
+		case ast.StateTypeChoice:
+			builder.WriteString(" : <<choice>>")
+		case ast.StateTypeHistory:
+			builder.WriteString(" : <<history>>")
+		case ast.StateTypeDeepHistory:
+			builder.WriteString(" : <<deepHistory>>")
+		}
+	}
+
+	builder.WriteString("\n")
+
+	// Render composite state body if it has sub-states
+	if len(state.SubStates) > 0 {
+		builder.WriteString("    state ")
+		builder.WriteString(state.ID)
+		builder.WriteString(" {\n")
+
+		for _, subState := range state.SubStates {
+			builder.WriteString("        state ")
+			builder.WriteString(subState.ID)
+			if subState.Label != subState.ID {
+				builder.WriteString(" as ")
+				builder.WriteString(subState.Label)
+			}
+			if subState.Description != nil {
+				builder.WriteString(" : ")
+				builder.WriteString(*subState.Description)
+			}
+			builder.WriteString("\n")
+		}
+
+		builder.WriteString("    }\n")
+	}
+
+	// Render note if present
+	if state.Note != nil {
+		builder.WriteString("    note ")
+		builder.WriteString(string(state.Note.Position))
+		builder.WriteString(" ")
+		builder.WriteString(state.ID)
+		builder.WriteString(" : ")
+		builder.WriteString(state.Note.Text)
+		builder.WriteString("\n")
+	}
+
+	// Render state actions
+	r.renderStateActions(builder, state)
+}
+
+// isStateInComposite checks if a state ID is within a composite state
+func (r *StateRenderer) isStateInComposite(compositeState *ast.StateNode, stateID string) bool {
+	// Don't consider the composite state itself as internal
+	if stateID == compositeState.ID {
+		return false
+	}
+
+	// Check if it's a direct substate
+	if _, exists := compositeState.SubStates[stateID]; exists {
+		return true
+	}
+
+	// Check if it's [*] (start/end states are considered internal to composite states)
+	if stateID == "[*]" {
+		return true
+	}
+
+	// Check if it's within any nested composite state
+	for _, subState := range compositeState.SubStates {
+		if len(subState.SubStates) > 0 {
+			// This is a nested composite state, check recursively
+			if r.isStateInComposite(subState, stateID) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// renderStateWithInternalTransitions renders a state with its internal transitions
+func (r *StateRenderer) renderStateWithInternalTransitions(builder *strings.Builder, state *ast.StateNode, internalTransitions []*ast.StateTransition) {
+	// For composite states, only render the composite structure
+	if len(state.SubStates) > 0 {
 		builder.WriteString("    state ")
 		builder.WriteString(state.ID)
 
@@ -67,31 +343,171 @@ func (r *StateRenderer) Render(diagram *ast.StateDiagram) (string, error) {
 			}
 		}
 
-		builder.WriteString("\n")
-
-		// Render composite state body if it has sub-states
-		if len(state.SubStates) > 0 {
-			builder.WriteString("    state ")
-			builder.WriteString(state.ID)
-			builder.WriteString(" {\n")
+		builder.WriteString(" {\n")
 
+		// Render transitions that belong to this composite state but not to nested composite states
+		for _, transition := range internalTransitions {
+			// Check if this transition belongs to a nested composite state
+			belongsToNested := false
 			for _, subState := range state.SubStates {
-				builder.WriteString("        state ")
-				builder.WriteString(subState.ID)
-				if subState.Label != subState.ID {
-					builder.WriteString(" as ")
-					builder.WriteString(subState.Label)
+				if len(subState.SubStates) > 0 {
+					// This is a nested composite state
+					if r.isStateInComposite(subState, transition.From) && r.isStateInComposite(subState, transition.To) {
+						belongsToNested = true
+						break
+					}
 				}
-				if subState.Description != nil {
+			}
+
+			if !belongsToNested {
+				builder.WriteString("        ")
+				builder.WriteString(transition.From)
+				builder.WriteString(" --> ")
+				builder.WriteString(transition.To)
+
+				// Add transition decorations
+				var decorations []string
+				if transition.Label != nil {
+					decorations = append(decorations, *transition.Label)
+				}
+				if transition.Condition != nil {
+					decorations = append(decorations, fmt.Sprintf("[%s]", *transition.Condition))
+				}
+				if transition.Action != nil {
+					decorations = append(decorations, fmt.Sprintf("/ %s", *transition.Action))
+				}
+
+				if len(decorations) > 0 {
 					builder.WriteString(" : ")
-					builder.WriteString(*subState.Description)
+					builder.WriteString(strings.Join(decorations, " "))
 				}
 				builder.WriteString("\n")
 			}
+		}
+
+		// Render nested composite states (in original order)
+		// Create ordered slice from map to ensure consistent ordering
+		var nestedStates []*ast.StateNode
+		for _, subState := range state.SubStates {
+			if len(subState.SubStates) > 0 {
+				nestedStates = append(nestedStates, subState)
+			}
+		}
+
+		// Sort nested states by their ID to ensure consistent ordering
+		for i := 0; i < len(nestedStates); i++ {
+			for j := i + 1; j < len(nestedStates); j++ {
+				if nestedStates[i].ID > nestedStates[j].ID {
+					nestedStates[i], nestedStates[j] = nestedStates[j], nestedStates[i]
+				}
+			}
+		}
+
+		for _, subState := range nestedStates {
+			// This is a nested composite state, render it
+			builder.WriteString("        state ")
+			builder.WriteString(subState.ID)
+
+			// Add alias if different from ID
+			if subState.Label != subState.ID {
+				builder.WriteString(" as ")
+				if strings.Contains(subState.Label, " ") {
+					builder.WriteString(fmt.Sprintf("\"%s\"", subState.Label))
+				} else {
+					builder.WriteString(subState.Label)
+				}
+			}
+
+			// Add description or special type
+			if subState.Description != nil {
+				builder.WriteString(" : ")
+				builder.WriteString(*subState.Description)
+			} else {
+				switch subState.Type {
+				case ast.StateTypeFork:
+					builder.WriteString(" : <<fork>>")
+				case ast.StateTypeJoin:
+					builder.WriteString(" : <<join>>")
+				case ast.StateTypeChoice:
+					builder.WriteString(" : <<choice>>")
+				case ast.StateTypeHistory:
+					builder.WriteString(" : <<history>>")
+				case ast.StateTypeDeepHistory:
+					builder.WriteString(" : <<deepHistory>>")
+				}
+			}
+
+			builder.WriteString(" {\n")
+
+			// Render nested state transitions
+			for _, nestedTransition := range internalTransitions {
+				// Check if this transition is within the nested state
+				if r.isStateInComposite(subState, nestedTransition.From) && r.isStateInComposite(subState, nestedTransition.To) {
+					builder.WriteString("            ")
+					builder.WriteString(nestedTransition.From)
+					builder.WriteString(" --> ")
+					builder.WriteString(nestedTransition.To)
+
+					// Add transition decorations
+					var decorations []string
+					if nestedTransition.Label != nil {
+						decorations = append(decorations, *nestedTransition.Label)
+					}
+					if nestedTransition.Condition != nil {
+						decorations = append(decorations, fmt.Sprintf("[%s]", *nestedTransition.Condition))
+					}
+					if nestedTransition.Action != nil {
+						decorations = append(decorations, fmt.Sprintf("/ %s", *nestedTransition.Action))
+					}
+
+					if len(decorations) > 0 {
+						builder.WriteString(" : ")
+						builder.WriteString(strings.Join(decorations, " "))
+					}
+					builder.WriteString("\n")
+				}
+			}
+
+			builder.WriteString("        }\n")
+		}
+
+		builder.WriteString("    }\n")
+	} else {
+		// For non-composite states, render the state declaration
+		builder.WriteString("    state ")
+		builder.WriteString(state.ID)
+
+		// Add alias if different from ID
+		if state.Label != state.ID {
+			builder.WriteString(" as ")
+			if strings.Contains(state.Label, " ") {
+				builder.WriteString(fmt.Sprintf("\"%s\"", state.Label))
+			} else {
+				builder.WriteString(state.Label)
+			}
+		}
 
-			builder.WriteString("    }\n")
+		// Add description or special type
+		if state.Description != nil {
+			builder.WriteString(" : ")
+			builder.WriteString(*state.Description)
+		} else {
+			switch state.Type {
+			case ast.StateTypeFork:
+				builder.WriteString(" : <<fork>>")
+			case ast.StateTypeJoin:
+				builder.WriteString(" : <<join>>")
+			case ast.StateTypeChoice:
+				builder.WriteString(" : <<choice>>")
+			case ast.StateTypeHistory:
+				builder.WriteString(" : <<history>>")
+			case ast.StateTypeDeepHistory:
+				builder.WriteString(" : <<deepHistory>>")
+			}
 		}
 
+		builder.WriteString("\n")
+
 		// Render note if present
 		if state.Note != nil {
 			builder.WriteString("    note ")
@@ -104,54 +520,6 @@ func (r *StateRenderer) Render(diagram *ast.StateDiagram) (string, error) {
 		}
 
 		// Render state actions
-		if state.EntryAction != nil {
-			builder.WriteString("    ")
-			builder.WriteString(state.ID)
-			builder.WriteString(" : entry ")
-			builder.WriteString(*state.EntryAction)
-			builder.WriteString("\n")
-		}
-		if state.ExitAction != nil {
-			builder.WriteString("    ")
-			builder.WriteString(state.ID)
-			builder.WriteString(" : exit ")
-			builder.WriteString(*state.ExitAction)
-			builder.WriteString("\n")
-		}
-		if state.DoAction != nil {
-			builder.WriteString("    ")
-			builder.WriteString(state.ID)
-			builder.WriteString(" : do ")
-			builder.WriteString(*state.DoAction)
-			builder.WriteString("\n")
-		}
+		r.renderStateActions(builder, state)
 	}
-
-	// Render transitions
-	for _, transition := range diagram.Transitions {
-		builder.WriteString("    ")
-		builder.WriteString(transition.From)
-		builder.WriteString(" --> ")
-		builder.WriteString(transition.To)
-
-		// Add transition decorations: label, [guard], /action
-		var decorations []string
-		if transition.Label != nil {
-			decorations = append(decorations, *transition.Label)
-		}
-		if transition.Condition != nil {
-			decorations = append(decorations, fmt.Sprintf("[%s]", *transition.Condition))
-		}
-		if transition.Action != nil {
-			decorations = append(decorations, fmt.Sprintf("/%s", *transition.Action))
-		}
-
-		if len(decorations) > 0 {
-			builder.WriteString(" : ")
-			builder.WriteString(strings.Join(decorations, " "))
-		}
-		builder.WriteString("\n")
-	}
-
-	return builder.String(), nil
 }