fix: preserve stratify drilldown state when navigating forward/backward

packages/client/hmi-client/src/components/models/tera-model-diagram.vue

@@ -45,8 +45,15 @@
 										"
 									/>
 								</span>
+								<Button
+									v-if="model && isStratifiedAMR(model) && !isEditing"
+									@click="toggleCollapsedView"
+									:label="isCollapsed ? 'Show expanded view' : 'Show collapsed view'"
+									class="p-button-sm p-button-outlined toolbar-button"
+								/>
 							</template>
 						</Toolbar>
+						<tera-model-type-legend v-if="model" :model="model" />
 						<div v-if="model" ref="graphElement" class="graph-element" />
 						<ContextMenu ref="menu" :model="contextMenuItems" />
 					</section>
@@ -226,6 +233,7 @@ import { Model, Observable } from '@/types/Types';
 import TeraModal from '@/components/widgets/tera-modal.vue';
 import InputText from 'primevue/inputtext';
 import TeraResizablePanel from '../widgets/tera-resizable-panel.vue';
+import TeraModelTypeLegend from './tera-model-type-legend.vue';
 
 // Get rid of these emits
 const emit = defineEmits([
@@ -457,8 +465,22 @@ const contextMenuItems = ref([
 	}
 ]);
 
+const isCollapsed = ref(true);
+async function toggleCollapsedView() {
+	isCollapsed.value = !isCollapsed.value;
+	if (props.model) {
+		const graphData: IGraph<NodeData, EdgeData> = convertToIGraphHelper(props.model);
+		// Render graph
+		if (renderer) {
+			renderer.isGraphDirty = true;
+			await renderer.setData(graphData);
+			await renderer.render();
+		}
+	}
+}
+
 const convertToIGraphHelper = (amr: Model) => {
-	if (isStratifiedAMR(amr)) {
+	if (isStratifiedAMR(amr) && isCollapsed.value) {
 		// FIXME: wont' work for MIRA
 		return convertToIGraph(props.model?.semantics?.span?.[0].system);
 	}
@@ -592,7 +614,7 @@ const cancelEdit = async () => {
 	if (!props.model) return;
 
 	// Convert petri net into a graph with raw input data
-	const graphData: IGraph<NodeData, EdgeData> = convertToIGraph(props.model);
+	const graphData: IGraph<NodeData, EdgeData> = convertToIGraphHelper(props.model);
 
 	if (renderer) {
 		renderer.setEditMode(false);

packages/client/hmi-client/src/components/models/tera-model-type-legend.vue

@@ -0,0 +1,86 @@
+<template>
+	<section class="legend" v-if="stateTypes || transitionTypes">
+		<ul>
+			<li v-for="(type, i) in stateTypes" :key="i">
+				<div class="legend-key-circle" :style="getLegendKeyStyle(type)" />
+				{{ type }}
+			</li>
+		</ul>
+		<ul>
+			<li v-for="(type, i) in transitionTypes" :key="i">
+				<div class="legend-key-square" :style="getLegendKeyStyle(type)" />
+				{{ type }}
+			</li>
+		</ul>
+	</section>
+</template>
+
+<script setup lang="ts">
+import { computed } from 'vue';
+import { Model } from '@/types/Types';
+import { useNodeTypeColorPalette } from '@/utils/petrinet-color-palette';
+
+const props = defineProps<{
+	model: Model;
+}>();
+
+const { getNodeTypeColor } = useNodeTypeColorPalette();
+
+const stateTypes = computed<string[]>(() =>
+	props.model.semantics?.typing?.system?.model.states.map((s) => s.name)
+);
+const transitionTypes = computed<string[]>(() =>
+	props.model.semantics?.typing?.system?.model.transitions.map((t) => t.properties?.name)
+);
+
+function getLegendKeyStyle(id: string) {
+	if (!id) {
+		return {
+			backgroundColor: 'var(--petri-nodeFill)'
+		};
+	}
+	return {
+		backgroundColor: getNodeTypeColor(id)
+	};
+}
+</script>
+
+<style scoped>
+.legend {
+	position: absolute;
+	bottom: 0;
+	z-index: 1;
+	margin-bottom: 1rem;
+	margin-left: 1rem;
+	display: flex;
+	gap: 1rem;
+	background-color: var(--surface-section);
+	border-radius: 0.5rem;
+	padding: 0.5rem;
+}
+
+.legend-key-circle {
+	height: 24px;
+	width: 24px;
+	border-radius: 12px;
+}
+
+.legend-key-square {
+	height: 24px;
+	width: 24px;
+	border-radius: 4px;
+}
+
+section.legend ul {
+	display: flex;
+	gap: 0.5rem;
+	list-style-type: none;
+	flex-direction: row;
+}
+
+section.legend li {
+	display: flex;
+	align-items: center;
+	gap: 0.5rem;
+}
+</style>

packages/client/hmi-client/src/components/models/tera-reflexives-toolbar.vue

@@ -12,7 +12,7 @@
 					optionLabel="id"
 					:model-value="statesToAddReflexives[transition.id]"
 					@update:model-value="
-						(newValue) => updateStatesToAddReflexives(newValue, transition, stateType as string)
+						(states) => updateStatesToAddReflexives({states, typeOfTransition: transition, typeIdOfState: stateType as string}, i)
 					"
 				/>
 			</div>
@@ -29,6 +29,7 @@ import {
 	addTyping,
 	updateRateExpression
 } from '@/model-representation/petrinet/petrinet-service';
+import { cloneDeep } from 'lodash';
 
 const props = defineProps<{
 	modelToUpdate: Model; // the model to which we will add reflexives
@@ -40,8 +41,7 @@ const emit = defineEmits(['model-updated']);
 const modelToCompareTypeSystem = computed<TypeSystem | undefined>(
 	() => props.modelToCompare.semantics?.typing?.system.model
 );
-const typedModel = ref<Model>(props.modelToUpdate); // this is the object that is being edited
-let unassignedTransitionTypes: Transition[] = [];
+const typedModel = ref<Model>(cloneDeep(props.modelToUpdate)); // this is the object that is being edited
 const statesToAddReflexives = ref<{ [id: string]: { id: string; name: string }[] }>({});
 const typeIdToTransitionIdMap = computed<{ [id: string]: string }>(() => {
 	const map: { [id: string]: string } = {};
@@ -63,7 +63,6 @@ const reflexiveNodeOptions = computed<{ [id: string]: { id: string; name: string
 	});
 	return options;
 });
-
 const stateId2NameMap = computed<{ [id: string]: string }>(() => {
 	const map: { [id: string]: string } = {};
 	props.modelToUpdate.model.states.forEach((state) => {
@@ -72,96 +71,118 @@ const stateId2NameMap = computed<{ [id: string]: string }>(() => {
 	return map;
 });
 
+const addedReflexivesRows: {
+	states: {
+		id: string;
+		name: string;
+	}[];
+	typeOfTransition: Transition;
+	typeIdOfState: string;
+}[] = [];
+
+/* Every time the user changes their selection of what states to add reflexives to, overwrite the previous changes with the current ones.
+   Since there can be multiple MultiSelect components, the selections for all MultiSelect components are combined in 'addedReflexivesRows'.
+   Iterate through 'addedReflexivesRows' and add reflexives according to the selections. 
+*/
 function updateStatesToAddReflexives(
-	newValue: {
-		id: string; // id of the state to which to add reflexive
-		name: string; // name of the state to which to add reflexive
-	}[],
-	typeOfTransition: Transition, // e.g. infect, recover
-	typeIdOfState: string // e.g. pop
+	selection: {
+		states: {
+			id: string;
+			name: string;
+		}[]; // list of id+name of state to which to add reflexives
+		typeOfTransition: Transition; // e.g. infect, recover
+		typeIdOfState: string; // e.g. pop
+	},
+	index: number
 ) {
-	statesToAddReflexives.value[typeOfTransition.id] = newValue;
-	const updatedTypeMap = typedModel.value.semantics?.typing?.map;
-	const updatedTypeSystem = typedModel.value.semantics?.typing?.system;
-
-	if (updatedTypeMap && updatedTypeSystem) {
-		newValue.forEach((state) => {
-			const newTransitionId = `${typeIdToTransitionIdMap.value[typeOfTransition.id]}${state.id}${
-				state.id
-			}`;
-			// For the type of reflexive transition that we are adding, get the number of inputs and outputs that share the same type as the state that we are updating
-			// E.g. if we are adding an 'Infect' reflexive to a node of type 'Pop', get the number of 'Pop' inputs and outputs for 'Infect'
-			const numInputsOfStateType = typeOfTransition.input.filter((i) => i === typeIdOfState).length;
-			// const numOutputsOfStateType = typeOfTransition.input.filter(i => i === typeIdOfState).length;
-			// Assume for now that the number of inputs and outputs for a given type are always equal, though in general this may not be the case
-			// TODO: implement logic for more generalized case where the above assumption is not true
-			addReflexives(typedModel.value, state.id, newTransitionId, numInputsOfStateType);
-			const reflexive = typedModel.value.model.transitions.find((t) => t.id === newTransitionId);
-
-			const transition = props.modelToCompare?.semantics?.typing?.system.model.transitions.find(
-				(t) => t.id === typeOfTransition.id
-			);
-			if (transition) {
-				updateRateExpression(typedModel.value, reflexive as PetriNetTransition, '');
-				if (!updatedTypeMap.find((m) => m[0] === newTransitionId)) {
-					updatedTypeMap.push([newTransitionId, typeOfTransition.id]);
+	typedModel.value = cloneDeep(props.modelToUpdate);
+	addedReflexivesRows[index] = selection;
+	addedReflexivesRows.forEach(({ states, typeOfTransition, typeIdOfState }) => {
+		statesToAddReflexives.value[typeOfTransition.id] = states;
+		const updatedTypeMap = typedModel.value.semantics?.typing?.map;
+		const updatedTypeSystem = typedModel.value.semantics?.typing?.system;
+		if (updatedTypeMap && updatedTypeSystem) {
+			states.forEach((state) => {
+				// For the type of reflexive transition that we are adding, get the number of inputs and outputs that share the same type as the state that we are updating
+				// E.g. if we are adding an 'Infect' reflexive to a node of type 'Pop', get the number of 'Pop' inputs and outputs for 'Infect'
+				const newTransitionId = `${typeIdToTransitionIdMap.value[typeOfTransition.id]}${state.id}`;
+
+				// Assume for now that the number of inputs and outputs for a given type are always equal, though in general this may not be the case
+				// TODO: implement logic for more generalized case where the above assumption is not true
+				const numInputsOfStateType = typeOfTransition.input.filter(
+					(i) => i === typeIdOfState
+				).length;
+				// const numOutputsOfStateType = typeOfTransition.input.filter(i => i === typeIdOfState).length;
+
+				if (!typedModel.value.model.transitions.find((t) => t.id === newTransitionId)) {
+					addReflexives(typedModel.value, state.id, newTransitionId, numInputsOfStateType);
 				}
-				if (!updatedTypeSystem.model.transitions.find((t) => t.id === typeOfTransition.id)) {
-					updatedTypeSystem.model.transitions.push(transition);
+				const reflexive = typedModel.value.model.transitions.find((t) => t.id === newTransitionId);
+
+				const transition = props.modelToCompare?.semantics?.typing?.system.model.transitions.find(
+					(t) => t.id === typeOfTransition.id
+				);
+				if (transition) {
+					updateRateExpression(typedModel.value, reflexive as PetriNetTransition, '');
+					if (!updatedTypeMap.find((m) => m[0] === newTransitionId)) {
+						updatedTypeMap.push([newTransitionId, typeOfTransition.id]);
+					}
+					if (!updatedTypeSystem.model.transitions.find((t) => t.id === typeOfTransition.id)) {
+						updatedTypeSystem.model.transitions.push(transition);
+					}
 				}
-			}
-		});
-		const updatedTyping: TypingSemantics = {
-			map: updatedTypeMap,
-			system: updatedTypeSystem
-		};
-		addTyping(typedModel.value, updatedTyping);
-	}
+			});
+			const updatedTyping: TypingSemantics = {
+				map: updatedTypeMap,
+				system: updatedTypeSystem
+			};
+			addTyping(typedModel.value, updatedTyping);
+		}
+	});
 	emit('model-updated', typedModel.value);
 }
 
-watch(
-	() => props.modelToUpdate,
-	() => {
-		if (props.modelToCompare) {
-			typedModel.value = props.modelToUpdate;
-			emit('model-updated', typedModel.value);
+/* Compare the type systems of 'modelToUpdate' and 'modelToCompare' to determine what options the user has for adding reflexives.
+   Allow the user to add transitions types that are present in 'modelToCompare' but not in 'modelToUpdate'.
+   E.g. if 'modelToUpdate' has ['Pop','Infect'] transitions and 'modelToCompare' has ['Pop,'Infect','Recover'] transitions, 
+   user should be prompted to add 'Recover' transitions to 'modelToUpdate'
+*/
+function populateReflexiveOptions() {
+	if (modelToCompareTypeSystem.value) {
+		let unassignedTransitions: Transition[];
+		const modelToUpdateTransitionIds =
+			props.modelToUpdate.semantics?.typing?.system.model.transitions.map((t) => t.id);
+		const modelToCompareTypeTransitionIds = modelToCompareTypeSystem.value?.transitions.map(
+			(t) => t.id
+		);
+		if (modelToUpdateTransitionIds && modelToCompareTypeTransitionIds) {
+			const unassignedIds = modelToCompareTypeTransitionIds.filter(
+				(id) => !modelToUpdateTransitionIds.includes(id)
+			);
+			// get the transition types that are in 'modelToCompare' but not 'modelToUpdate'
+			unassignedTransitions = modelToCompareTypeSystem.value?.transitions.filter((t) =>
+				unassignedIds.includes(t.id)
+			);
 		}
-	},
-	{ immediate: true }
-);
+		props.modelToUpdate.model.states.forEach((state) => {
+			// get type of state for each state in model to update model
+			const type: string =
+				props.modelToUpdate.semantics?.typing?.map.find((m) => m[0] === state.id)?.[1] ?? '';
+			// for each unassigned transition type, check if inputs or ouputs have the type of this state
+			// you should only be allowed to add a transition to a state, if the transition has inputs or outputs of the same type as the state
+			const allowedTransitionsForState: Transition[] = unassignedTransitions.filter(
+				(unassigned) => unassigned.input.includes(type) || unassigned.output.includes(type)
+			);
+			reflexiveOptions.value[type] = allowedTransitionsForState;
+		});
+	}
+}
 
 watch(
-	[() => modelToCompareTypeSystem],
+	[() => props.modelToCompare, () => props.modelToUpdate.semantics?.typing],
 	() => {
-		if (modelToCompareTypeSystem.value) {
-			const modelToUpdateTransitionIds =
-				props.modelToUpdate.semantics?.typing?.system.model.transitions.map((t) => t.id);
-			const modelToCompareTypeTransitionIds = modelToCompareTypeSystem.value?.transitions.map(
-				(t) => t.id
-			);
-			if (modelToUpdateTransitionIds && modelToCompareTypeTransitionIds) {
-				const unassignedIds = modelToCompareTypeTransitionIds.filter(
-					(id) => !modelToUpdateTransitionIds.includes(id)
-				);
-				const unassignedTransitions: Transition[] =
-					modelToCompareTypeSystem.value?.transitions.filter((t) => unassignedIds.includes(t.id));
-				if (unassignedTransitions.length > 0) {
-					unassignedTransitionTypes = unassignedTransitionTypes.concat(unassignedTransitions);
-				}
-			}
-			props.modelToUpdate.model.states.forEach((state) => {
-				// get type of state for each state in model to update model
-				const type: string =
-					props.modelToUpdate.semantics?.typing?.map.find((m) => m[0] === state.id)?.[1] ?? '';
-				// for each unassigned transition type, check if inputs or ouputs have the type of this state
-				const allowedTransitionsForState: Transition[] = unassignedTransitionTypes.filter(
-					(unassigned) => unassigned.input.includes(type) || unassigned.output.includes(type)
-				);
-				if (!reflexiveOptions.value[type]) {
-					reflexiveOptions.value[type] = allowedTransitionsForState;
-				}
-			});
+		if (props.modelToUpdate.semantics?.typing) {
+			populateReflexiveOptions();
 		}
 	},
 	{ immediate: true }