From 06f220d6a9a3561e20cb5c658319c1be8715fc9a Mon Sep 17 00:00:00 2001
From: SiboVG <sibo.vangool@hotmail.com>
Date: Wed, 2 Mar 2022 22:39:20 +0100
Subject: [PATCH] [#1208] Add StageSelector button for ParallelStage

---
 .../rocketcomponent/RocketComponent.java      | 51 ++++++++++++++++++-
 .../gui/components/StageSelector.java         |  9 +++-
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/core/src/net/sf/openrocket/rocketcomponent/RocketComponent.java b/core/src/net/sf/openrocket/rocketcomponent/RocketComponent.java
index 266daec44..54a62b6e0 100644
--- a/core/src/net/sf/openrocket/rocketcomponent/RocketComponent.java
+++ b/core/src/net/sf/openrocket/rocketcomponent/RocketComponent.java
@@ -1588,12 +1588,38 @@ public abstract class RocketComponent implements ChangeSource, Cloneable, Iterab
 		this.checkComponentStructure();
 		return children.get(n);
 	}
-	
+
+	/**
+	 * Returns all the direct children of this component. The result is a clone of the children list and may be edited.
+	 * @return direct children of this component.
+	 */
 	public final List<RocketComponent> getChildren() {
 		checkState();
 		this.checkComponentStructure();
 		return children.clone();
 	}
+
+	/**
+	 * Returns all children of this component, as well as the children of the children etc.
+	 * @return all the children (direct and indirect) of this component
+	 */
+	public final List<RocketComponent> getChildrenRecursive() {
+		checkState();
+		List<RocketComponent> children = getChildren();
+
+		if (children == null) {
+			return null;
+		}
+
+		for (RocketComponent child : new ArrayList<>(children)) {
+			List<RocketComponent> temp = child.getChildrenRecursive();
+			if (temp != null && temp.size() > 0) {
+				children.addAll(temp);
+			}
+		}
+
+		return children;
+	}
 	
 	
 	/**
@@ -1683,10 +1709,31 @@ public abstract class RocketComponent implements ChangeSource, Cloneable, Iterab
 		while ( null != curComponent ) {
 			if( ComponentAssembly.class.isAssignableFrom( curComponent.getClass()))
 				return (ComponentAssembly) curComponent;
-			curComponent = curComponent.parent;
 		}
 		throw new IllegalStateException("getAssembly() called on hierarchy without a ComponentAssembly.");
 	}
+
+	/**
+	 * Return all the component assemblies that are a child of this component
+	 * @return list of ComponentAssembly components that are a child of this component
+	 */
+	public final List<RocketComponent> getChildAssemblies() {
+		checkState();
+
+		List<RocketComponent> children = getChildrenRecursive();
+		if (children == null) {
+			return null;
+		}
+
+		List<RocketComponent> result = new ArrayList<>();
+
+		for (RocketComponent child : children) {
+			if (child instanceof ComponentAssembly) {
+				result.add(child);
+			}
+		}
+		return result;
+	}
 	
 	
 	/**
diff --git a/swing/src/net/sf/openrocket/gui/components/StageSelector.java b/swing/src/net/sf/openrocket/gui/components/StageSelector.java
index 2b36a940e..3dd580855 100644
--- a/swing/src/net/sf/openrocket/gui/components/StageSelector.java
+++ b/swing/src/net/sf/openrocket/gui/components/StageSelector.java
@@ -12,6 +12,7 @@ import javax.swing.JToggleButton;
 import net.miginfocom.swing.MigLayout;
 import net.sf.openrocket.gui.widgets.SelectColorToggleButton;
 import net.sf.openrocket.rocketcomponent.AxialStage;
+import net.sf.openrocket.rocketcomponent.BodyTube;
 import net.sf.openrocket.rocketcomponent.ComponentChangeEvent;
 import net.sf.openrocket.rocketcomponent.FlightConfiguration;
 import net.sf.openrocket.rocketcomponent.Rocket;
@@ -36,7 +37,11 @@ public class StageSelector extends JPanel implements StateChangeListener {
 	private void updateButtons( final FlightConfiguration configuration ) {
 		buttons.clear();
 		this.removeAll();
-		for(RocketComponent stage : configuration.getRocket().getChildren()){
+		List<RocketComponent> assemblies = configuration.getRocket().getChildAssemblies();
+		if (assemblies == null) {
+			return;
+		}
+		for (RocketComponent stage : assemblies) {
 			if (!(stage instanceof AxialStage)) continue;
 			JToggleButton button = new SelectColorToggleButton(new StageAction((AxialStage) stage));
 			button.setSelected(configuration.isStageActive(stage.getStageNumber()));
@@ -50,7 +55,7 @@ public class StageSelector extends JPanel implements StateChangeListener {
 	@Override
 	public void stateChanged(EventObject eo) {
 		Object source = eo.getSource();
-		if ((source instanceof Rocket) || (source instanceof AxialStage)) {
+		if ((source instanceof Rocket) || (source instanceof AxialStage) || (source instanceof BodyTube)) {
 			Rocket rkt = (Rocket) ((RocketComponent) source).getRoot();
 			updateButtons( rkt.getSelectedConfiguration() );
 		}