egui_treeize/ui/
state.rs

1use egui::{
2  Context, Id, Pos2, Rect, Ui, Vec2,
3  ahash::HashSet,
4  emath::{GuiRounding, TSTransform},
5  style::Spacing,
6};
7use smallvec::{SmallVec, ToSmallVec, smallvec};
8
9use crate::{InPinId, NodeId, OutPinId, Treeize};
10
11use super::{TreeizeWidget, transform_matching_points};
12
13pub type RowHeights = SmallVec<[f32; 8]>;
14
15/// Node UI state.
16#[derive(Debug)]
17pub struct NodeState {
18  /// Node size for this frame.
19  /// It is updated to fit content.
20  size: Vec2,
21  header_height: f32,
22  input_heights: RowHeights,
23  output_heights: RowHeights,
24
25  id: Id,
26  dirty: bool,
27}
28
29#[derive(Clone, PartialEq)]
30pub(crate) struct NodeData {
31  pub(crate) size: Vec2,
32  header_height: f32,
33  input_heights: RowHeights,
34  output_heights: RowHeights,
35}
36
37impl NodeState {
38  pub fn load(cx: &Context, id: Id, spacing: &Spacing) -> Self {
39    cx.data(|d| d.get_temp::<NodeData>(id)).map_or_else(
40      || {
41        cx.request_discard("NodeState initialization");
42        Self::initial(id, spacing)
43      },
44      |data| NodeState {
45        size: data.size,
46        header_height: data.header_height,
47        input_heights: data.input_heights,
48        output_heights: data.output_heights,
49        id,
50        dirty: false,
51      },
52    )
53  }
54
55  pub fn pick_data(cx: &Context, id: Id) -> Option<NodeData> {
56    cx.data(|d| d.get_temp::<NodeData>(id))
57  }
58
59  pub fn clear(self, cx: &Context) {
60    cx.data_mut(|d| d.remove::<Self>(self.id));
61  }
62
63  pub fn store(self, cx: &Context) {
64    if self.dirty {
65      cx.data_mut(|d| {
66        d.insert_temp(
67          self.id,
68          NodeData {
69            size: self.size,
70            header_height: self.header_height,
71            input_heights: self.input_heights,
72            output_heights: self.output_heights,
73          },
74        );
75      });
76      cx.request_repaint();
77    }
78  }
79
80  /// Finds node rect at specific position (excluding node frame margin).
81  pub fn node_rect(&self, pos: Pos2, openness: f32) -> Rect {
82    Rect::from_min_size(
83      pos,
84      egui::vec2(self.size.x, f32::max(self.header_height, self.size.y * openness)),
85    )
86    .round_ui()
87  }
88
89  pub fn payload_offset(&self, openness: f32) -> f32 {
90    ((self.size.y) * (1.0 - openness)).round_ui()
91  }
92
93  pub fn set_size(&mut self, size: Vec2) {
94    if self.size != size {
95      self.size = size;
96      self.dirty = true;
97    }
98  }
99
100  pub fn header_height(&self) -> f32 {
101    self.header_height.round_ui()
102  }
103
104  pub fn set_header_height(&mut self, height: f32) {
105    #[allow(clippy::float_cmp)]
106    if self.header_height != height {
107      self.header_height = height;
108      self.dirty = true;
109    }
110  }
111
112  pub const fn input_heights(&self) -> &RowHeights {
113    &self.input_heights
114  }
115
116  pub const fn output_heights(&self) -> &RowHeights {
117    &self.output_heights
118  }
119
120  const fn initial(id: Id, spacing: &Spacing) -> Self {
121    NodeState {
122      size: spacing.interact_size,
123      header_height: spacing.interact_size.y,
124      input_heights: SmallVec::new_const(),
125      output_heights: SmallVec::new_const(),
126      id,
127      dirty: true,
128    }
129  }
130}
131
132#[derive(Clone)]
133pub enum NewWires {
134  In(SmallVec<[InPinId; 4]>),
135  Out(SmallVec<[OutPinId; 4]>),
136}
137
138#[derive(Clone, Copy)]
139struct RectSelect {
140  origin: Pos2,
141  current: Pos2,
142}
143
144pub struct TreeizeState {
145  /// Treeize viewport transform to global space.
146  to_global: TSTransform,
147
148  new_wires: Option<NewWires>,
149
150  /// Flag indicating that new wires are owned by the menu now.
151  new_wires_menu: bool,
152
153  id: Id,
154
155  /// Flag indicating that the graph state is dirty must be saved.
156  dirty: bool,
157
158  /// Active rect selection.
159  rect_selection: Option<RectSelect>,
160
161  /// Order of nodes to draw.
162  draw_order: Vec<NodeId>,
163
164  /// List of currently selected nodes.
165  selected_nodes: SmallVec<[NodeId; 8]>,
166}
167
168#[derive(Clone, Default)]
169struct DrawOrder(Vec<NodeId>);
170
171impl DrawOrder {
172  fn save(self, cx: &Context, id: Id) {
173    cx.data_mut(|d| {
174      if self.0.is_empty() {
175        d.remove_temp::<Self>(id);
176      } else {
177        d.insert_temp::<Self>(id, self);
178      }
179    });
180  }
181
182  fn load(cx: &Context, id: Id) -> Self {
183    cx.data(|d| d.get_temp::<Self>(id)).unwrap_or_default()
184  }
185}
186
187#[derive(Clone, Default)]
188struct SelectedNodes(SmallVec<[NodeId; 8]>);
189
190impl SelectedNodes {
191  fn save(self, cx: &Context, id: Id) {
192    cx.data_mut(|d| {
193      if self.0.is_empty() {
194        d.remove_temp::<Self>(id);
195      } else {
196        d.get_temp_mut_or_default::<Self>(id).clone_from(&self);
197        d.insert_temp::<Self>(id, self);
198      }
199    });
200  }
201
202  fn load(cx: &Context, id: Id) -> Self {
203    cx.data(|d| d.get_temp::<Self>(id)).unwrap_or_default()
204  }
205}
206
207#[derive(Clone)]
208struct TreeizeStateData {
209  to_global: TSTransform,
210  new_wires: Option<NewWires>,
211  new_wires_menu: bool,
212  rect_selection: Option<RectSelect>,
213}
214
215impl TreeizeStateData {
216  fn save(self, cx: &Context, id: Id) {
217    cx.data_mut(|d| {
218      d.insert_temp(id, self);
219    });
220  }
221
222  fn load(cx: &Context, id: Id) -> Option<Self> {
223    cx.data(|d| d.get_temp(id))
224  }
225}
226
227fn prune_selected_nodes<T>(
228  selected_nodes: &mut SmallVec<[NodeId; 8]>,
229  treeize: &Treeize<T>,
230) -> bool {
231  let old_size = selected_nodes.len();
232  selected_nodes.retain(|node| treeize.nodes.contains(node.0));
233  old_size != selected_nodes.len()
234}
235
236impl TreeizeState {
237  pub fn load<T>(
238    cx: &Context,
239    id: Id,
240    treeize: &Treeize<T>,
241    ui_rect: Rect,
242    min_scale: f32,
243    max_scale: f32,
244  ) -> Self {
245    let Some(data) = TreeizeStateData::load(cx, id) else {
246      cx.request_discard("Initial placing");
247      return Self::initial(id, treeize, ui_rect, min_scale, max_scale);
248    };
249
250    let mut selected_nodes = SelectedNodes::load(cx, id).0;
251    let dirty = prune_selected_nodes(&mut selected_nodes, treeize);
252
253    let draw_order = DrawOrder::load(cx, id).0;
254
255    TreeizeState {
256      to_global: data.to_global,
257      new_wires: data.new_wires,
258      new_wires_menu: data.new_wires_menu,
259      id,
260      dirty,
261      rect_selection: data.rect_selection,
262      draw_order,
263      selected_nodes,
264    }
265  }
266
267  fn initial<T>(
268    id: Id,
269    treeize: &Treeize<T>,
270    ui_rect: Rect,
271    min_scale: f32,
272    max_scale: f32,
273  ) -> Self {
274    let mut bb = Rect::NOTHING;
275
276    for (_, node) in &treeize.nodes {
277      bb.extend_with(node.pos);
278    }
279
280    if bb.is_finite() {
281      bb = bb.expand(100.0);
282    } else if ui_rect.is_finite() {
283      bb = ui_rect;
284    } else {
285      bb = Rect::from_min_max(Pos2::new(-100.0, -100.0), Pos2::new(100.0, 100.0));
286    }
287
288    let scaling2 = ui_rect.size() / bb.size();
289    let scaling = scaling2.min_elem().clamp(min_scale, max_scale);
290
291    let to_global = transform_matching_points(bb.center(), ui_rect.center(), scaling);
292
293    TreeizeState {
294      to_global,
295      new_wires: None,
296      new_wires_menu: false,
297      id,
298      dirty: true,
299      draw_order: Vec::new(),
300      rect_selection: None,
301      selected_nodes: SmallVec::new(),
302    }
303  }
304
305  #[inline(always)]
306  pub fn store<T>(mut self, treeize: &Treeize<T>, cx: &Context) {
307    self.dirty |= prune_selected_nodes(&mut self.selected_nodes, treeize);
308
309    if self.dirty {
310      let data = TreeizeStateData {
311        to_global: self.to_global,
312        new_wires: self.new_wires,
313        new_wires_menu: self.new_wires_menu,
314        rect_selection: self.rect_selection,
315      };
316      data.save(cx, self.id);
317
318      DrawOrder(self.draw_order).save(cx, self.id);
319      SelectedNodes(self.selected_nodes).save(cx, self.id);
320
321      cx.request_repaint();
322    }
323  }
324
325  pub const fn to_global(&self) -> TSTransform {
326    self.to_global
327  }
328
329  pub fn set_to_global(&mut self, to_global: TSTransform) {
330    if self.to_global != to_global {
331      self.to_global = to_global;
332      self.dirty = true;
333    }
334  }
335
336  pub fn look_at(&mut self, view: Rect, ui_rect: Rect, min_scale: f32, max_scale: f32) {
337    let scaling2 = ui_rect.size() / view.size();
338    let scaling = scaling2.min_elem().clamp(min_scale, max_scale);
339
340    let to_global = transform_matching_points(view.center(), ui_rect.center(), scaling);
341
342    if self.to_global != to_global {
343      self.to_global = to_global;
344      self.dirty = true;
345    }
346  }
347
348  pub fn start_new_wire_in(&mut self, pin: InPinId) {
349    self.new_wires = Some(NewWires::In(smallvec![pin]));
350    self.new_wires_menu = false;
351    self.dirty = true;
352  }
353
354  pub fn start_new_wire_out(&mut self, pin: OutPinId) {
355    self.new_wires = Some(NewWires::Out(smallvec![pin]));
356    self.new_wires_menu = false;
357    self.dirty = true;
358  }
359
360  pub fn start_new_wires_in(&mut self, pins: &[InPinId]) {
361    self.new_wires = Some(NewWires::In(pins.to_smallvec()));
362    self.new_wires_menu = false;
363    self.dirty = true;
364  }
365
366  pub fn start_new_wires_out(&mut self, pins: &[OutPinId]) {
367    self.new_wires = Some(NewWires::Out(pins.to_smallvec()));
368    self.new_wires_menu = false;
369    self.dirty = true;
370  }
371
372  pub fn add_new_wire_in(&mut self, pin: InPinId) {
373    debug_assert!(!self.new_wires_menu);
374    let Some(NewWires::In(pins)) = &mut self.new_wires else {
375      unreachable!();
376    };
377
378    if !pins.contains(&pin) {
379      pins.push(pin);
380      self.dirty = true;
381    }
382  }
383
384  pub fn add_new_wire_out(&mut self, pin: OutPinId) {
385    debug_assert!(!self.new_wires_menu);
386    let Some(NewWires::Out(pins)) = &mut self.new_wires else {
387      unreachable!();
388    };
389
390    if !pins.contains(&pin) {
391      pins.push(pin);
392      self.dirty = true;
393    }
394  }
395
396  pub fn remove_new_wire_in(&mut self, pin: InPinId) {
397    debug_assert!(!self.new_wires_menu);
398    let Some(NewWires::In(pins)) = &mut self.new_wires else {
399      unreachable!();
400    };
401
402    if let Some(idx) = pins.iter().position(|p| *p == pin) {
403      pins.swap_remove(idx);
404      self.dirty = true;
405    }
406  }
407
408  pub fn remove_new_wire_out(&mut self, pin: OutPinId) {
409    debug_assert!(!self.new_wires_menu);
410    let Some(NewWires::Out(pins)) = &mut self.new_wires else {
411      unreachable!();
412    };
413
414    if let Some(idx) = pins.iter().position(|p| *p == pin) {
415      pins.swap_remove(idx);
416      self.dirty = true;
417    }
418  }
419
420  pub const fn has_new_wires(&self) -> bool {
421    matches!((self.new_wires.as_ref(), self.new_wires_menu), (Some(_), false))
422  }
423
424  pub const fn has_new_wires_in(&self) -> bool {
425    matches!((&self.new_wires, self.new_wires_menu), (Some(NewWires::In(_)), false))
426  }
427
428  pub const fn has_new_wires_out(&self) -> bool {
429    matches!((&self.new_wires, self.new_wires_menu), (Some(NewWires::Out(_)), false))
430  }
431
432  pub const fn new_wires(&self) -> Option<&NewWires> {
433    match (&self.new_wires, self.new_wires_menu) {
434      (Some(new_wires), false) => Some(new_wires),
435      _ => None,
436    }
437  }
438
439  pub const fn take_new_wires(&mut self) -> Option<NewWires> {
440    match (&self.new_wires, self.new_wires_menu) {
441      (Some(_), false) => {
442        self.dirty = true;
443        self.new_wires.take()
444      }
445      _ => None,
446    }
447  }
448
449  pub(crate) const fn take_new_wires_menu(&mut self) -> Option<NewWires> {
450    match (&self.new_wires, self.new_wires_menu) {
451      (Some(_), true) => {
452        self.dirty = true;
453        self.new_wires.take()
454      }
455      _ => None,
456    }
457  }
458
459  pub(crate) fn set_new_wires_menu(&mut self, wires: NewWires) {
460    debug_assert!(self.new_wires.is_none());
461    self.new_wires = Some(wires);
462    self.new_wires_menu = true;
463  }
464
465  pub(crate) fn update_draw_order<T>(&mut self, treeize: &Treeize<T>) -> Vec<NodeId> {
466    let mut node_ids = treeize.nodes.iter().map(|(id, _)| NodeId(id)).collect::<HashSet<_>>();
467
468    self.draw_order.retain(|id| {
469      let has = node_ids.remove(id);
470      self.dirty |= !has;
471      has
472    });
473
474    self.dirty |= !node_ids.is_empty();
475
476    for new_id in node_ids {
477      self.draw_order.push(new_id);
478    }
479
480    self.draw_order.clone()
481  }
482
483  pub(crate) fn node_to_top(&mut self, node: NodeId) {
484    if let Some(order) = self.draw_order.iter().position(|idx| *idx == node) {
485      self.draw_order.remove(order);
486      self.draw_order.push(node);
487    }
488    self.dirty = true;
489  }
490
491  pub fn selected_nodes(&self) -> &[NodeId] {
492    &self.selected_nodes
493  }
494
495  pub fn select_one_node(&mut self, reset: bool, node: NodeId) {
496    if reset {
497      if self.selected_nodes[..] == [node] {
498        return;
499      }
500
501      self.deselect_all_nodes();
502    } else if let Some(pos) = self.selected_nodes.iter().position(|n| *n == node) {
503      if pos == self.selected_nodes.len() - 1 {
504        return;
505      }
506      self.selected_nodes.remove(pos);
507    }
508    self.selected_nodes.push(node);
509    self.dirty = true;
510  }
511
512  pub fn select_many_nodes(&mut self, reset: bool, nodes: impl Iterator<Item = NodeId>) {
513    if reset {
514      self.deselect_all_nodes();
515      self.selected_nodes.extend(nodes);
516      self.dirty = true;
517    } else {
518      nodes.for_each(|node| self.select_one_node(false, node));
519    }
520  }
521
522  pub fn deselect_one_node(&mut self, node: NodeId) {
523    if let Some(pos) = self.selected_nodes.iter().position(|n| *n == node) {
524      self.selected_nodes.remove(pos);
525      self.dirty = true;
526    }
527  }
528
529  pub fn deselect_many_nodes(&mut self, nodes: impl Iterator<Item = NodeId>) {
530    for node in nodes {
531      if let Some(pos) = self.selected_nodes.iter().position(|n| *n == node) {
532        self.selected_nodes.remove(pos);
533        self.dirty = true;
534      }
535    }
536  }
537
538  pub fn deselect_all_nodes(&mut self) {
539    self.dirty |= !self.selected_nodes.is_empty();
540    self.selected_nodes.clear();
541  }
542
543  pub const fn start_rect_selection(&mut self, pos: Pos2) {
544    self.dirty |= self.rect_selection.is_none();
545    self.rect_selection = Some(RectSelect { origin: pos, current: pos });
546  }
547
548  pub const fn stop_rect_selection(&mut self) {
549    self.dirty |= self.rect_selection.is_some();
550    self.rect_selection = None;
551  }
552
553  pub const fn is_rect_selection(&self) -> bool {
554    self.rect_selection.is_some()
555  }
556
557  pub const fn update_rect_selection(&mut self, pos: Pos2) {
558    if let Some(rect_selection) = &mut self.rect_selection {
559      rect_selection.current = pos;
560      self.dirty = true;
561    }
562  }
563
564  pub fn rect_selection(&self) -> Option<Rect> {
565    let rect = self.rect_selection?;
566    Some(Rect::from_two_pos(rect.origin, rect.current))
567  }
568}
569
570impl TreeizeWidget {
571  /// Returns list of nodes selected in the UI for the `TreeizeWidget` with same id.
572  ///
573  /// Use same `Ui` instance that was used in [`TreeizeWidget::show`].
574  #[must_use]
575  #[inline]
576  pub fn get_selected_nodes(self, ui: &Ui) -> Vec<NodeId> {
577    self.get_selected_nodes_at(ui.id(), ui.ctx())
578  }
579
580  /// Returns list of nodes selected in the UI for the `TreeizeWidget` with same id.
581  ///
582  /// `ui_id` must be the Id of the `Ui` instance that was used in [`TreeizeWidget::show`].
583  #[must_use]
584  #[inline]
585  pub fn get_selected_nodes_at(self, ui_id: Id, ctx: &Context) -> Vec<NodeId> {
586    let treeize_id = self.get_id(ui_id);
587
588    ctx.data(|d| d.get_temp::<SelectedNodes>(treeize_id).unwrap_or_default().0).into_vec()
589  }
590}
591
592/// Returns nodes selected in the UI for the `TreeizeWidget` with same ID.
593///
594/// Only works if [`TreeizeWidget::id`] was used.
595/// For other cases construct [`TreeizeWidget`] and use [`TreeizeWidget::get_selected_nodes`] or [`TreeizeWidget::get_selected_nodes_at`].
596#[must_use]
597#[inline]
598pub fn get_selected_nodes(id: Id, ctx: &Context) -> Vec<NodeId> {
599  ctx.data(|d| d.get_temp::<SelectedNodes>(id).unwrap_or_default().0).into_vec()
600}
egui_treeize/ui/state.rs

egui_treeize/ui/
state.rs
