diff --git a/src/core/setup.typ b/src/core/setup.typ
new file mode 100644
index 0000000..e8d07e9
--- /dev/null
+++ b/src/core/setup.typ
@@ -0,0 +1,317 @@
+#import "utils.typ": is-elmt, get-group-span
+#import "../participant.typ": _exists as par-exists, _par
+
+#let flatten-group(elmts, i) = {
+  let group = elmts.at(i)
+  elmts.at(i) = group
+  return (
+    elmts.slice(0, i + 1) +
+    group.elmts +
+    ((
+      type: "grp-end",
+      start-i: i
+    ),) +
+    elmts.slice(i+1)
+  )
+}
+
+#let update-group-children(elmts, i) = {
+  let elmts = elmts
+  let group-end = elmts.at(i)
+  
+  elmts.at(group-end.start-i).elmts = elmts.slice(group-end.start-i + 1, i)
+  return elmts
+}
+
+#let convert-return(elmts, i, activation-history) = {
+  if activation-history.len() == 0 {
+    panic("Cannot return if no lifeline is activated")
+  }
+  let elmts = elmts
+  let activation-history = activation-history
+  let ret = elmts.at(i)
+  let seq = activation-history.pop()
+  elmts.at(i) = _seq(
+    seq.p2, seq.p1,
+    comment: ret.comment,
+    disable-src: true,
+    dashed: true
+  ).first()
+  return (elmts, activation-history)
+}
+
+#let unwrap-containers(elmts) = {
+  let elmts = elmts
+  let i = 0
+  let activation-history = ()
+
+  // Flatten groups + convert returns
+  while i < elmts.len() {
+    let elmt = elmts.at(i)
+    if not is-elmt(elmt) {
+      i += 1
+      continue
+    }
+
+    if elmt.type == "grp" {
+      elmts = flatten-group(elmts, i)
+
+    } else if elmt.type == "seq" {
+      if elmt.enable-dst {
+        activation-history.push(elmt)
+      }
+    
+    } else if elmt.type == "evt" {
+      if elmt.event == "enable" {
+        for elmt2 in elmts.slice(0, i).rev() {
+          if elmt2.type == "seq" {
+            activation-history.push(elmt2)
+            break
+          }
+        }
+      }
+    
+    } else if elmt.type == "ret" {
+      (elmts, activation-history) = convert-return(elmts, i, activation-history)
+    }
+    i += 1
+  }
+
+  return (elmts, activation-history)
+}
+
+
+#let prepare-seq-participants(ctx, seq) = {
+  let ctx = ctx
+  if not par-exists(ctx.participants, seq.p1) {
+    ctx.participants.push(_par(seq.p1).first())
+  }
+  if not par-exists(ctx.participants, seq.p2) {
+    ctx.participants.push(_par(
+      seq.p2,
+      from-start: not seq.create-dst
+    ).first())
+  
+  } else if seq.create-dst {
+    let i = ctx.participants.position(p => p.name == seq.p2)
+    ctx.participants.at(i).from-start = false
+  }
+
+  let p1 = seq.p1
+  let p2 = seq.p2
+  if seq.p1 == "?" {
+    p1 = "?" + seq.p2
+  }
+  if seq.p2 == "?" {
+    p2 = seq.p1 + "?"
+  }
+  ctx.linked.push(p1)
+  ctx.linked.push(p2)
+  ctx.last-seq = (
+    elmt: seq,
+    i: ctx.i,
+    p1: p1,
+    p2: p2
+  )
+  return ctx
+}
+
+#let prepare-note-participants(ctx, note) = {
+  let ctx = ctx
+  let note = note
+  note.insert(
+    "linked",
+    note.pos == none and note.side != "across"
+  )
+  if note.pos == none and note.side != "across" {
+    let names = ctx.participants.map(p => p.name)
+    let i1 = names.position(n => n == ctx.last-seq.p1)
+    let i2 = names.position(n => n == ctx.last-seq.p2)
+    let pars = (
+      (i1, ctx.last-seq.p1),
+      (i2, ctx.last-seq.p2)
+    ).sorted(key: p => p.first())
+
+    if note.side == "left" {
+      note.pos = pars.first().last()
+    } else if note.side == "right" {
+      note.pos = pars.last().last()
+    }
+
+    let seq = last-seq.note
+    seq.insert("linked-note", note)
+    ctx.elmts.at(last-seq.i) = seq
+  }
+  if note.aligned {
+    let n = last-note.note
+    n.aligned-with = note
+    ctx.elmts.at(last-note.i) = n
+  }
+  if note.side == "left" {
+    ctx.linked.push("[")
+  } else if note.side == "right" {
+    ctx.linked.push("]")
+  }
+
+  let pars = none
+  if type(note.pos) == str {
+    pars = (note.pos,)
+  } else if type(note.pos) == array {
+    pars = note.pos
+  }
+  if pars != none {
+    for par in pars {
+      if not par-exists(ctx.participants, par) {
+        participants.push(_par(par).first())
+      }
+    }
+  }
+
+  ctx.elmts.at(ctx.i) = note
+
+  ctx.last-note = (
+    elmt: note,
+    i: ctx.i
+  )
+
+  return ctx
+}
+
+#let prepare-evt-participants(ctx, evt) = {
+  let par = evt.participant
+  if not par-exists(ctx.participants, par) {
+    let p = _par(
+      par,
+      from-start: evt.event != "create"
+    ).first()
+    ctx.participants.push(p)
+  
+  } else if evt.event == "create" {
+    let i = ctx.participants.position(p => p.name == par)
+    ctx.participants.at(i).from-start = false
+  }
+  return ctx
+}
+
+#let normalize-special-participants(elmt) = {
+  if elmt.p1 == "?" {
+    elmt.p1 = "?" + elmt.p2
+  } else if elmt.p2 == "?" {
+    elmt.p2 = elmt.p1 + "?"
+  }
+  return elmt
+}
+
+#let prepare-participants(elmts) = {
+  let ctx = (
+    linked: (),
+    last-seq: none,
+    last-note: none,
+    participants: (),
+    elmts: elmts,
+    i: 0
+  )
+
+  for (i, elmt) in ctx.elmts.enumerate() {
+    ctx.i = i
+    if not is-elmt(elmt) {
+      continue
+    }
+
+    if elmt.type == "par" {
+      ctx.participants.push(elmt)
+
+    } else if elmt.type == "seq" {
+      ctx = prepare-seq-participants(ctx, elmt)
+    
+    } else if elmt.type == "note" {
+      ctx = prepare-note-participants(ctx, elmt)
+    
+    } else if elmt.type == "evt" {
+      ctx = prepare-evt-participants(ctx, elmt)
+    }
+  }
+  ctx.linked = ctx.linked.dedup()
+
+  let pars = ctx.participants
+  let participants = ()
+  
+  if "[" in ctx.linked {
+    participants.push(_par("[", invisible: true).first())
+  }
+  
+  for (i, p) in pars.enumerate() {
+    let before = _par(
+      "?" + p.name,
+      invisible: true
+    ).first()
+    let after = _par(
+      p.name + "?",
+      invisible: true
+    ).first()
+    
+    if before.name in ctx.linked {
+      if participants.len() == 0 or not participants.last().name.ends-with("?") {
+        participants.push(before)
+      } else {
+        participants.insert(-1, before)
+      }
+    }
+
+    participants.push(p)
+
+    if after.name in ctx.linked {
+      participants.push(after)
+    }
+  }
+  if "]" in ctx.linked {
+    participants.push(_par(
+      "]",
+      invisible: true
+    ).first())
+  }
+
+  return (ctx.elmts, participants)
+}
+
+#let finalize-setup(elmts, participants) = {
+  for (i, p) in participants.enumerate() {
+    p.insert("i", i)
+    participants.at(i) = p
+  }
+
+  let containers = ()
+
+  for (i, elmt) in elmts.enumerate() {
+    if not is-elmt(elmt) {
+      continue
+    }
+    if elmt.type == "seq" {
+      elmts.at(i) = normalize-special-participants(elmt)
+    } else if elmt.type == "grp-end" {
+      // Put back elements in group because they might have changed
+      elmts = update-group-children(elmts, i)
+    } else if elmt.type in ("grp", "alt") {
+      containers.push(i)
+    }
+  }
+
+  // Compute groups spans (horizontal)
+  for i in containers {
+    let elmt = elmts.at(i)
+    let (min-i, max-i) = get-group-span(participants, elmt)
+    elmts.at(i).insert("min-i", min-i)
+    elmts.at(i).insert("max-i", max-i)
+  }
+
+  return (elmts, participants)
+}
+
+#let setup(elements) = {
+  let (elmts, activation-history) = unwrap-containers(elements)
+  
+  let participants
+  (elmts, participants) = prepare-participants(elmts)
+
+  return finalize-setup(elmts, participants)
+}
\ No newline at end of file
diff --git a/src/utils.typ b/src/core/utils.typ
similarity index 93%
rename from src/utils.typ
rename to src/core/utils.typ
index e8d750d..25d4801 100644
--- a/src/utils.typ
+++ b/src/core/utils.typ
@@ -1,3 +1,13 @@
+#let is-elmt(elmt) = {
+  if type(elmt) != dictionary {
+    return false
+  }
+  if "type" not in elmt {
+    return false
+  }
+  return true
+}
+
 #let normalize-units(value) = {
   if type(value) == int or type(value) == float {
     return value
diff --git a/src/diagram.typ b/src/diagram.typ
index 0e78117..0dc89aa 100644
--- a/src/diagram.typ
+++ b/src/diagram.typ
@@ -1,8 +1,10 @@
-#import "utils.typ": get-group-span, fit-canvas
+#import "core/utils.typ": fit-canvas
 #import "renderer.typ": render
 #import "participant.typ" as participant: _par, PAR-SPECIALS
 #import "sequence.typ": _seq
 
+#import "core/setup.typ": setup
+
 #let _gap(size: 20) = {
   return ((
     type: "gap",
@@ -36,214 +38,7 @@
     return
   }
   
-  let participants = ()
-  let elmts = elements
-  let i = 0
-
-  let activation-history = ()
-
-  // Flatten groups + convert returns
-  while i < elmts.len() {
-    let elmt = elmts.at(i)
-    if elmt.type == "grp" {
-      elmt.elmts = elmt.elmts.map(e => {
-        if e.type == "seq" {
-          if e.p1 == "?" {
-            e.p1 = "?" + e.p2
-          } else if e.p2 == "?" {
-            e.p2 = e.p1 + "?"
-          }
-        }
-        e
-      })
-      elmts.at(i) = elmt
-      elmts = (
-        elmts.slice(0, i + 1) +
-        elmt.elmts +
-        ((
-          type: "grp-end",
-          start-i: i
-        ),) +
-        elmts.slice(i+1)
-      )
-    } else if elmt.type == "grp-end" {
-      // Put back elements in group because they might have changed
-      elmts.at(elmt.start-i).elmts = elmts.slice(elmt.start-i + 1, i)
-
-    } else if elmt.type == "seq" {
-      if elmt.enable-dst {
-        activation-history.push(elmt)
-      }
-    } else if elmt.type == "evt" {
-      if elmt.event == "enable" {
-        for elmt2 in elmts.slice(0, i).rev() {
-          if elmt2.type == "seq" {
-            activation-history.push(elmt2)
-            break
-          }
-        }
-      }
-    } else if elmt.type == "ret" {
-      if activation-history.len() == 0 {
-        panic("Cannot return if no lifeline is activated")
-      }
-      let seq = activation-history.pop()
-      elmts.at(i) = _seq(
-        seq.p2, seq.p1,
-        comment: elmt.comment,
-        disable-src: true,
-        dashed: true
-      ).first()
-    }
-    i += 1
-  }
-
-  // List participants
-  let linked = ()
-  let last-seq = none
-  let last-note = none
-  for (i, elmt) in elmts.enumerate() {
-    if elmt.type == "par" {
-      participants.push(elmt)
-    } else if elmt.type == "seq" {
-      if not participant._exists(participants, elmt.p1) {
-        participants.push(_par(elmt.p1).first())
-      }
-      if not participant._exists(participants, elmt.p2) {
-        let par = _par(elmt.p2, from-start: not elmt.create-dst).first()
-        participants.push(par)
-      
-      } else if elmt.create-dst {
-        let i = participants.position(p => p.name == elmt.p2)
-        participants.at(i).from-start = false
-      }
-
-      let p1 = elmt.p1
-      let p2 = elmt.p2
-      if elmt.p1 == "?" {
-        p1 = "?" + elmt.p2
-      }
-      if elmt.p2 == "?" {
-        p2 = elmt.p1 + "?"
-      }
-      linked.push(p1)
-      linked.push(p2)
-      last-seq = (
-        elmt: elmt,
-        i: i,
-        p1: p1,
-        p2: p2
-      )
-    } else if elmt.type == "note" {
-      elmt.insert("linked", elmt.pos == none and elmt.side != "across")
-      if elmt.pos == none and elmt.side != "across" {
-        let names = participants.map(p => p.name)
-        let i1 = names.position(n => n == last-seq.p1)
-        let i2 = names.position(n => n == last-seq.p2)
-        let pars = ((i1, last-seq.p1), (i2, last-seq.p2)).sorted(key: p => p.first())
-        if elmt.side == "left" {
-          elmt.pos = pars.first().last()
-        } else if elmt.side == "right" {
-          elmt.pos = pars.last().last()
-        }
-
-        let seq = last-seq.elmt
-        seq.insert("linked-note", elmt)
-        elmts.at(last-seq.i) = seq
-      }
-      if elmt.aligned {
-        let n = last-note.elmt
-        n.aligned-with = elmt
-        elmts.at(last-note.i) = n
-      }
-      elmts.at(i) = elmt
-      if elmt.side == "left" {
-        linked.push("[")
-      } else if elmt.side == "right" {
-        linked.push("]")
-      }
-
-      let pars = none
-      if type(elmt.pos) == str {
-        pars = (elmt.pos,)
-      } else if type(elmt.pos) == array {
-        pars = elmt.pos
-      }
-      if pars != none {
-        for par in pars {
-          if not participant._exists(participants, par) {
-            participants.push(_par(par).first())
-          }
-        }
-      }
-
-      last-note = (
-        elmt: elmt,
-        i: i
-      )
-    } else if elmt.type == "evt" {
-      let par = elmt.participant
-      if not participant._exists(participants, par) {
-        let p = _par(par, from-start: elmt.event != "create").first()
-        participants.push(p)
-      
-      } else if elmt.event == "create" {
-        let i = participants.position(p => p.name == par)
-        participants.at(i).from-start = false
-      }
-    }
-  }
-  linked = linked.dedup()
-
-  let pars = participants
-  participants = ()
-  
-  if "[" in linked {
-    participants.push(_par("[", invisible: true).first())
-  }
-  
-  for (i, p) in pars.enumerate() {
-    let before = _par("?" + p.name, invisible: true).first()
-    let after = _par(p.name + "?", invisible: true).first()
-    
-    if before.name in linked {
-      if participants.len() == 0 or not participants.last().name.ends-with("?") {
-        participants.push(before)
-      } else {
-        participants.insert(-1, before)
-      }
-    }
-
-    participants.push(p)
-
-    if after.name in linked {
-      participants.push(after)
-    }
-  }
-  if "]" in linked {
-    participants.push(_par("]", invisible: true).first())
-  }
-
-  // Add index to participant
-  for (i, p) in participants.enumerate() {
-    p.insert("i", i)
-    participants.at(i) = p
-  }
-
-  // Compute groups spans (horizontal)
-  for (i, elmt) in elmts.enumerate() {
-    if elmt.type == "grp" or elmt.type == "alt" {
-      let (min-i, max-i) = get-group-span(participants, elmt)
-      elmts.at(i).insert("min-i", min-i)
-      elmts.at(i).insert("max-i", max-i)
-    } else if elmt.type == "seq" {
-      if elmt.p1 == "?" {
-        elmts.at(i).p1 = "?" + elmt.p2
-      } else if elmt.p2 == "?" {
-        elmts.at(i).p2 = elmt.p1 + "?"
-      }
-    }
-  }
+  let (elmts, participants) = setup(elements)
 
   let canvas = render(participants, elmts)
   fit-canvas(canvas, width: width)