diagram 0.2.1 → 0.3.2

data/lib/diagrams/elements/timeline_event.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Diagrams
+  module Elements
+    # Represents a single event description within a timeline period.
+    class TimelineEvent < Dry::Struct
+      include Elements::Types
+
+      attribute :description, Types::Strict::String.constrained(min_size: 1)
+
+      # Returns a hash representation suitable for serialization.
+      #
+      # @return [Hash{Symbol => String}]
+      def to_h
+        {
+          description:
+        }
+      end
+    end
+  end
+end

data/lib/diagrams/elements/timeline_period.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Diagrams
+  module Elements
+    # Represents a specific time period on the timeline, containing one or more events.
+    class TimelinePeriod < Dry::Struct
+      include Elements::Types
+
+      attribute :label, Types::Strict::String.constrained(min_size: 1)
+      attribute :events, Types::Strict::Array.of(TimelineEvent).constrained(min_size: 1)
+
+      # Returns a hash representation suitable for serialization.
+      #
+      # @return [Hash{Symbol => String | Array<Hash>}]
+      def to_h
+        {
+          label:,
+          events: events.map(&:to_h)
+        }
+      end
+    end
+  end
+end

data/lib/diagrams/elements/timeline_section.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Diagrams
+  module Elements
+    # Represents a section or age within the timeline, grouping multiple time periods.
+    class TimelineSection < Dry::Struct
+      include Elements::Types
+
+      attribute :title, Types::Strict::String.constrained(min_size: 1)
+      attribute :periods, Types::Strict::Array.of(TimelinePeriod).default([].freeze)
+
+      # Returns a hash representation suitable for serialization.
+      #
+      # @return [Hash{Symbol => String | Array<Hash>}]
+      def to_h
+        {
+          title:,
+          periods: periods.map(&:to_h)
+        }
+      end
+    end
+  end
+end

data/lib/diagrams/elements/transition.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Diagrams
+  module Elements
+    # Represents a transition between two states in a State Diagram.
+    class Transition < Dry::Struct
+      # Use the shared Types module
+      include Elements::Types
+
+      # Consider if a transition needs its own ID.
+      # attribute :id, Types::Strict::String
+
+      attribute :source_state_id, Types::Strict::String.constrained(min_size: 1)
+      attribute :target_state_id, Types::Strict::String.constrained(min_size: 1)
+      # Label often represents the event or condition triggering the transition.
+      attribute :label, Types::Strict::String.optional.default(nil)
+
+      # Returns a hash representation suitable for serialization.
+      #
+      # @return [Hash{Symbol => String | nil}]
+      def to_h
+        # Rely on Dry::Struct's default to_h, filtering out nil label.
+        super.compact
+      end
+    end
+  end
+end

data/lib/diagrams/elements.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Diagrams
+  # Namespace for diagram element value objects.
+
+  module Elements
+    # Common Dry::Types definitions for elements.
+    module Types
+      include Dry.Types()
+    end
+  end
+end

data/lib/diagrams/flowchart_diagram.rb

@@ -1,9 +1,124 @@
 # frozen_string_literal: true
 
 module Diagrams
-  class FlowchartDiagram < AbstractDiagram
-    attribute :id, Types::String
-    attribute :nodes, Types::Array.of(Node)
-    attribute :links, Types::Array.of(Link)
+  # Represents a flowchart diagram consisting of nodes and edges connecting them.
+  class FlowchartDiagram < Base
+    attr_reader :nodes, :edges
+
+    # Initializes a new FlowchartDiagram.
+    #
+    # @param nodes [Array<Element::Node>] An array of node objects.
+    # @param edges [Array<Element::Edge>] An array of edge objects.
+    # @param version [String, Integer, nil] User-defined version identifier.
+    def initialize(nodes: [], edges: [], version: 1)
+      super(version:)
+      @nodes = nodes.is_a?(Array) ? nodes : []
+      @edges = edges.is_a?(Array) ? edges : []
+      validate_elements!
+      update_checksum!
+    end
+
+    # Adds a node to the diagram.
+    #
+    # @param node [Element::Node] The node object to add.
+    # @raise [ArgumentError] if a node with the same ID already exists.
+    # @return [Element::Node] The added node.
+    def add_node(node)
+      raise ArgumentError, 'Node must be a Diagrams::Element::Node' unless node.is_a?(Diagrams::Elements::Node)
+      raise ArgumentError, "Node with ID '#{node.id}' already exists" if find_node(node.id)
+
+      @nodes << node
+      update_checksum!
+      node
+    end
+
+    # Adds an edge to the diagram.
+    #
+    # @param edge [Element::Edge] The edge object to add.
+    # @raise [ArgumentError] if the edge refers to non-existent node IDs.
+    # @return [Element::Edge] The added edge.
+    def add_edge(edge)
+      raise ArgumentError, 'Edge must be a Diagrams::Element::Edge' unless edge.is_a?(Diagrams::Elements::Edge)
+      unless find_node(edge.source_id) && find_node(edge.target_id)
+        raise ArgumentError, "Edge refers to non-existent node IDs ('#{edge.source_id}' or '#{edge.target_id}')"
+      end
+
+      @edges << edge
+      update_checksum!
+      edge
+    end
+
+    # Finds a node by its ID.
+    #
+    # @param node_id [String] The ID of the node to find.
+    # @return [Element::Node, nil] The found node or nil.
+    def find_node(node_id)
+      @nodes.find { |n| n.id == node_id }
+    end
+
+    # Returns the specific content of the flowchart diagram as a hash.
+    # Called by `Diagrams::Base#to_h`.
+    #
+    # @return [Hash{Symbol => Array<Hash>}]
+    def to_h_content
+      {
+        nodes: @nodes.map(&:to_h),
+        edges: @edges.map(&:to_h)
+      }
+    end
+
+    # Returns a hash mapping element types to their collections for diffing.
+    # @see Diagrams::Base#identifiable_elements
+    # @return [Hash{Symbol => Array<Diagrams::Elements::Node | Diagrams::Elements::Edge>}]
+    def identifiable_elements
+      {
+        nodes: @nodes,
+        edges: @edges
+      }
+    end
+
+    # Class method to create a FlowchartDiagram from a hash.
+    # Used by the deserialization factory in `Diagrams::Base`.
+    #
+    # @param data_hash [Hash] Hash containing `:nodes` and `:edges` arrays.
+    # @param version [String, Integer, nil] Diagram version.
+    # @param checksum [String, nil] Expected checksum (optional, for verification).
+    # @return [FlowchartDiagram] The instantiated diagram.
+    def self.from_h(data_hash, version:, checksum:)
+      nodes_data = data_hash[:nodes] || data_hash['nodes'] || []
+      edges_data = data_hash[:edges] || data_hash['edges'] || []
+
+      nodes =
+        nodes_data.map do |node_h|
+          Diagrams::Elements::Node.new(node_h.transform_keys(&:to_sym))
+        end
+      edges =
+        edges_data.map do |edge_h|
+          Diagrams::Elements::Edge.new(edge_h.transform_keys(&:to_sym))
+        end
+      diagram = new(nodes:, edges:, version:)
+
+      # Optional: Verify checksum if provided
+      if checksum && diagram.checksum != checksum
+        warn "Checksum mismatch for loaded FlowchartDiagram (version: #{version}). Expected #{checksum}, got #{diagram.checksum}."
+        # Or raise an error: raise "Checksum mismatch..."
+      end
+
+      diagram
+    end
+
+    private
+
+    # Validates the consistency of nodes and edges during initialization.
+    def validate_elements!
+      node_ids = @nodes.map(&:id)
+      raise ArgumentError, 'Duplicate node IDs found' unless node_ids.uniq.size == @nodes.size
+
+      @edges.each do |edge|
+        unless node_ids.include?(edge.source_id) && node_ids.include?(edge.target_id)
+          raise ArgumentError, "Edge refers to non-existent node IDs ('#{edge.source_id}' or '#{edge.target_id}')"
+        end
+      end
+    end
   end
 end

data/lib/diagrams/gantt_diagram.rb

@@ -1,8 +1,99 @@
 # frozen_string_literal: true
 
 module Diagrams
-  class GanttDiagram < AbstractDiagram
-    attribute :title, GanttDiagram::Types::String
-    attribute :sections, GanttDiagram::Types::Array.of(Section)
+  # Represents a Gantt Chart diagram consisting of tasks over time.
+  class GanttDiagram < Base
+    attr_reader :title, :tasks
+
+    # Initializes a new GanttDiagram.
+    #
+    # @param title [String] The title of the Gantt chart.
+    # @param tasks [Array<Element::Task>] An array of task objects.
+    # @param version [String, Integer, nil] User-defined version identifier.
+    def initialize(title: '', tasks: [], version: 1)
+      super(version:)
+      @title = title.is_a?(String) ? title : ''
+      @tasks = tasks.is_a?(Array) ? tasks : []
+      validate_elements!
+      update_checksum!
+    end
+
+    # Adds a task to the diagram.
+    #
+    # @param task [Element::Task] The task object to add.
+    # @raise [ArgumentError] if a task with the same ID already exists.
+    # @return [Element::Task] The added task.
+    def add_task(task)
+      raise ArgumentError, 'Task must be a Diagrams::Elements::Task' unless task.is_a?(Diagrams::Elements::Task)
+      raise ArgumentError, "Task with ID '#{task.id}' already exists" if find_task(task.id)
+
+      @tasks << task
+      update_checksum!
+      task
+    end
+
+    # Finds a task by its ID.
+    #
+    # @param task_id [String] The ID of the task to find.
+    # @return [Element::Task, nil] The found task or nil.
+    def find_task(task_id)
+      @tasks.find { |t| t.id == task_id }
+    end
+
+    # Returns the specific content of the Gantt diagram as a hash.
+    # Called by `Diagrams::Base#to_h`.
+    #
+    # @return [Hash{Symbol => String | Array<Hash>}]
+    def to_h_content
+      {
+        title: @title,
+        tasks: @tasks.map(&:to_h)
+      }
+    end
+
+    # Returns a hash mapping element types to their collections for diffing.
+    # @see Diagrams::Base#identifiable_elements
+    # @return [Hash{Symbol => Array<Diagrams::Elements::Task>}]
+    def identifiable_elements
+      {
+        tasks: @tasks
+      }
+    end
+
+    # Class method to create a GanttDiagram from a hash.
+    # Used by the deserialization factory in `Diagrams::Base`.
+    #
+    # @param data_hash [Hash] Hash containing `:title` and `:tasks` array.
+    # @param version [String, Integer, nil] Diagram version.
+    # @param checksum [String, nil] Expected checksum (optional, for verification).
+    # @return [GanttDiagram] The instantiated diagram.
+    def self.from_h(data_hash, version:, checksum:)
+      title = data_hash[:title] || data_hash['title'] || ''
+      tasks_data = data_hash[:tasks] || data_hash['tasks'] || []
+
+      tasks = tasks_data.map { |task_h| Diagrams::Elements::Task.new(task_h.transform_keys(&:to_sym)) }
+
+      diagram = new(title:, tasks:, version:)
+
+      # Optional: Verify checksum if provided
+      if checksum && diagram.checksum != checksum
+        warn "Checksum mismatch for loaded GanttDiagram (version: #{version}). Expected #{checksum}, got #{diagram.checksum}."
+        # Or raise an error: raise "Checksum mismatch..."
+      end
+
+      diagram
+    end
+
+    private
+
+    # Validates the consistency of tasks during initialization.
+    def validate_elements!
+      task_ids = @tasks.map(&:id)
+      return if task_ids.uniq.size == @tasks.size
+
+      raise ArgumentError, 'Duplicate task IDs found'
+
+      # Add more validation if needed (e.g., date formats, dependencies)
+    end
   end
 end

data/lib/diagrams/gitgraph_diagram.rb

@@ -0,0 +1,345 @@
+# frozen_string_literal: true
+
+require 'digest' # For generating default commit IDs if needed
+
+module Diagrams
+  # Represents a Gitgraph diagram, tracking commits, branches, and their relationships.
+  class GitgraphDiagram < Base
+    attr_reader :commits, :branches, :commit_order, :current_branch_name
+
+    # Initializes a new GitgraphDiagram.
+    # Starts with a 'master' branch by default.
+    #
+    # @param version [String, Integer, nil] User-defined version identifier.
+    def initialize(version: 1)
+      super
+      @commits = {} # Hash { commit_id => GitCommit }
+      @branches = {} # Hash { branch_name => GitBranch }
+      @commit_order = [] # Array<String> - IDs of commits in order of creation/operation
+      @current_branch_name = 'master'
+
+      # Initialize main branch conceptually. Its start/head commit will be set by the first commit.
+      # We need a placeholder start_commit_id; using a special value or handling nil in GitBranch might be better.
+      # For now, let's use a placeholder that signifies it's the root.
+      # A better approach might be to create the branch *during* the first commit. Let's refine this.
+      # --> Refinement: Don't create the branch object here. Create it during the first 'commit' or 'branch' operation.
+      #                 Initialize @current_branch_name = 'master' conceptually.
+
+      update_checksum! # Initial checksum for an empty graph
+    end
+
+    # --- Git Operations ---
+
+    # Adds a commit to the current branch.
+    # Handles the creation of the initial 'master' branch on the first commit.
+    #
+    # @param id [String, nil] Optional custom ID for the commit. Auto-generated if nil.
+    # @param message [String, nil] Optional commit message.
+    # @param tag [String, nil] Optional tag for the commit.
+    # @param type [Symbol] Type of the commit (:NORMAL, :REVERSE, :HIGHLIGHT).
+    # @raise [ArgumentError] if a commit with the given ID already exists.
+    # @return [Elements::GitCommit] The created commit object.
+    def commit(id: nil, message: nil, tag: nil, type: :NORMAL)
+      parent_id = current_head_commit_id
+      parent_ids = parent_id ? [parent_id] : []
+
+      commit_id = id || generate_commit_id(parent_ids, message)
+      raise ArgumentError, "Commit with ID '#{commit_id}' already exists" if @commits.key?(commit_id)
+
+      # Handle first commit: create the master branch
+      if @commits.empty? && @current_branch_name == 'master' && !@branches.key?('master')
+        # The first commit *is* the starting point of the master branch
+        master_branch = Elements::GitBranch.new(name: 'master', start_commit_id: commit_id, head_commit_id: commit_id)
+        @branches['master'] = master_branch
+      elsif !@branches.key?(@current_branch_name)
+        # This case shouldn't typically happen if branch/checkout is used correctly,
+        # but defensively handle committing to a non-existent branch (other than initial master).
+        raise ArgumentError, "Cannot commit: Branch '#{@current_branch_name}' does not exist."
+      end
+
+      new_commit = Elements::GitCommit.new(
+        id: commit_id,
+        parent_ids:,
+        branch_name: @current_branch_name,
+        message:,
+        tag:,
+        type:
+      )
+
+      @commits[commit_id] = new_commit
+      @commit_order << commit_id
+
+      # Update the head of the current branch
+      current_branch = @branches[@current_branch_name]
+      current_branch.attributes[:head_commit_id] = commit_id # Update using Dry::Struct's way if needed, direct assign might work
+
+      update_checksum!
+      new_commit
+    end
+
+    # Creates a new branch pointing to a specific commit (or the current head)
+    # and switches the current context to the new branch.
+    #
+    # @param name [String] The name for the new branch.
+    # @param start_commit_id [String, nil] Optional ID of the commit where the branch should start.
+    #                                     Defaults to the head commit of the current branch.
+    # @raise [ArgumentError] if the branch name already exists or if trying to branch before any commits exist.
+    # @raise [ArgumentError] if a specified `start_commit_id` does not exist.
+    # @return [Elements::GitBranch] The created branch object.
+    def branch(name:, start_commit_id: nil)
+      raise ArgumentError, "Branch name '#{name}' already exists" if @branches.key?(name)
+
+      effective_start_commit_id = start_commit_id || current_head_commit_id
+
+      # Ensure there's a commit to branch from
+      raise ArgumentError, 'Cannot create a branch before the first commit' unless effective_start_commit_id
+
+      unless @commits.key?(effective_start_commit_id)
+        raise ArgumentError,
+              "Start commit ID '#{effective_start_commit_id}' does not exist"
+      end
+
+      new_branch = Elements::GitBranch.new(
+        name:,
+        # The new branch initially points to the commit it was created from
+        start_commit_id: effective_start_commit_id,
+        head_commit_id: effective_start_commit_id
+      )
+
+      @branches[name] = new_branch
+      @current_branch_name = name # Switch to the new branch
+
+      update_checksum!
+      new_branch
+    end
+
+    # Switches the current context to an existing branch.
+    #
+    # @param name [String] The name of the branch to switch to.
+    # @raise [ArgumentError] if the branch name does not exist.
+    # @return [String] The name of the branch checked out.
+    def checkout(name:)
+      raise ArgumentError, "Branch '#{name}' does not exist. Cannot checkout." unless @branches.key?(name)
+
+      @current_branch_name = name
+      # NOTE: Checkout does not change the diagram structure itself (commits/branches),
+      # so we do NOT update the checksum here.
+      name
+    end
+
+    # Merges the head of a specified branch into the current branch.
+    # Creates a merge commit on the current branch.
+    #
+    # @param from_branch_name [String] The name of the branch to merge from.
+    # @param id [String, nil] Optional custom ID for the merge commit. Auto-generated if nil.
+    # @param tag [String, nil] Optional tag for the merge commit.
+    # @param type [Symbol] Type of the merge commit (defaults to :MERGE, can be overridden e.g., :REVERSE).
+    # @raise [ArgumentError] if `from_branch_name` does not exist, is the same as the current branch,
+    #                        or if either branch has no commits.
+    # @raise [ArgumentError] if a commit with the given ID already exists.
+    # @return [Elements::GitCommit] The created merge commit object.
+    def merge(from_branch_name:, id: nil, tag: nil, type: :MERGE)
+      if from_branch_name == @current_branch_name
+        raise ArgumentError,
+              "Cannot merge branch '#{from_branch_name}' into itself"
+      end
+      unless @branches.key?(from_branch_name)
+        raise ArgumentError,
+              "Branch '#{from_branch_name}' does not exist. Cannot merge."
+      end
+      unless @branches.key?(@current_branch_name)
+        raise ArgumentError, "Current branch '#{@current_branch_name}' does not exist. Cannot merge."
+      end
+
+      target_branch = @branches[@current_branch_name]
+      source_branch = @branches[from_branch_name]
+
+      target_head_id = target_branch.head_commit_id
+      source_head_id = source_branch.head_commit_id
+
+      unless target_head_id
+        raise ArgumentError,
+              "Current branch '#{@current_branch_name}' has no commits to merge into."
+      end
+      raise ArgumentError, "Source branch '#{from_branch_name}' has no commits to merge from." unless source_head_id
+
+      # Merge commit parents are the heads of the two branches being merged
+      parent_ids = [target_head_id, source_head_id].sort # Sort for consistent checksumming/comparison
+
+      merge_commit_id = id || generate_commit_id(parent_ids,
+                                                 "Merge branch '#{from_branch_name}' into #{@current_branch_name}")
+      raise ArgumentError, "Commit with ID '#{merge_commit_id}' already exists" if @commits.key?(merge_commit_id)
+
+      merge_commit = Elements::GitCommit.new(
+        id: merge_commit_id,
+        parent_ids:,
+        branch_name: @current_branch_name, # Merge commit belongs to the target branch
+        message: "Merge branch '#{from_branch_name}' into #{@current_branch_name}", # Default message
+        tag:,
+        type: # Use provided type, default :MERGE
+      )
+
+      @commits[merge_commit_id] = merge_commit
+      @commit_order << merge_commit_id
+
+      # Update the head of the current (target) branch
+      target_branch.attributes[:head_commit_id] = merge_commit_id
+
+      update_checksum!
+      merge_commit
+    end
+
+    # Cherry-picks an existing commit onto the current branch.
+    # Creates a new commit on the current branch that mirrors the specified commit.
+    #
+    # @param commit_id [String] The ID of the commit to cherry-pick.
+    # @param parent_override_id [String, nil] Optional: If cherry-picking a merge commit, specifies which parent lineage to follow.
+    #                                         (Note: Basic implementation might ignore this for simplicity initially).
+    # @raise [ArgumentError] if the commit_id does not exist, is already on the current branch,
+    #                        or if the current branch has no commits.
+    # @return [Elements::GitCommit] The created cherry-pick commit object.
+    # Basic implementation ignores parent_override_id for now
+    def cherry_pick(commit_id:, parent_override_id: nil)
+      unless @commits.key?(commit_id)
+        raise ArgumentError,
+              "Commit with ID '#{commit_id}' does not exist. Cannot cherry-pick."
+      end
+
+      source_commit = @commits[commit_id]
+      current_branch_head_id = current_head_commit_id
+
+      unless current_branch_head_id
+        raise ArgumentError,
+              "Current branch '#{@current_branch_name}' has no commits. Cannot cherry-pick onto it."
+      end
+      if source_commit.branch_name == @current_branch_name
+        raise ArgumentError,
+              "Commit '#{commit_id}' is already on the current branch '#{@current_branch_name}'. Cannot cherry-pick."
+      end
+
+      # More robust check: walk history? For now, simple branch name check.
+
+      # TODO: Handle cherry-picking merge commits and parent_override_id if needed later.
+      if source_commit.parent_ids.length > 1 && !parent_override_id
+        warn "Cherry-picking a merge commit (#{commit_id}) without specifying a parent override is ambiguous. Picking first parent lineage by default."
+        # Or raise ArgumentError: "Cherry-picking a merge commit requires specifying parent_override_id."
+      end
+
+      parent_ids = [current_branch_head_id] # Cherry-pick commit's parent is the current head
+      new_commit_id = generate_commit_id(parent_ids, "Cherry-pick: #{source_commit.message || source_commit.id}")
+      if @commits.key?(new_commit_id)
+        raise ArgumentError,
+              "Generated commit ID '#{new_commit_id}' conflicts with existing commit."
+      end
+
+      cherry_pick_commit = Elements::GitCommit.new(
+        id: new_commit_id,
+        parent_ids:,
+        branch_name: @current_branch_name,
+        message: source_commit.message || "Cherry-pick of #{source_commit.id}", # Copy message or use default
+        tag: nil, # Cherry-picks usually don't copy tags directly
+        type: :CHERRY_PICK,
+        cherry_pick_source_id: commit_id # Link back to the original commit
+      )
+
+      @commits[new_commit_id] = cherry_pick_commit
+      @commit_order << new_commit_id
+
+      # Update the head of the current branch
+      current_branch = @branches[@current_branch_name]
+      current_branch.attributes[:head_commit_id] = new_commit_id
+
+      update_checksum!
+      cherry_pick_commit
+    end
+
+    # --- Base Class Implementation ---
+
+    # Returns the specific content of the gitgraph diagram as a hash.
+    # @return [Hash]
+    def to_h_content
+      {
+        commits: @commits.values.map(&:to_h),
+        branches: @branches.values.map(&:to_h),
+        commit_order: @commit_order,
+        current_branch_name: @current_branch_name # Useful for resuming state? Maybe not needed in content hash.
+        # Consider if current_branch_name should be part of the checksummable content.
+        # For now, let's include it for potential deserialization needs.
+      }
+    end
+
+    # Returns a hash mapping element types to their collections for diffing.
+    # @return [Hash{Symbol => Array<Elements::GitCommit | Elements::GitBranch>}]
+    def identifiable_elements
+      {
+        commits: @commits.values,
+        branches: @branches.values
+      }
+    end
+
+    # Class method to create a GitgraphDiagram from a hash.
+    # @param data_hash [Hash] Hash containing diagram data.
+    # @param version [String, Integer, nil] Diagram version.
+    # @param checksum [String, nil] Expected checksum (optional).
+    # @return [GitgraphDiagram] The instantiated diagram.
+    def self.from_h(data_hash, version:, checksum:)
+      diagram = new(version:)
+
+      # Restore commits
+      commits_data = data_hash[:commits] || data_hash['commits'] || []
+      commits_data.each do |commit_h|
+        # Convert type back to symbol if it's a string
+        commit_data = commit_h.transform_keys(&:to_sym)
+        commit_data[:type] = commit_data[:type].to_sym if commit_data[:type].is_a?(String)
+        commit = Elements::GitCommit.new(commit_data)
+        diagram.commits[commit.id] = commit
+      end
+
+      # Restore branches
+      branches_data = data_hash[:branches] || data_hash['branches'] || []
+      branches_data.each do |branch_h|
+        branch = Elements::GitBranch.new(branch_h.transform_keys(&:to_sym))
+        diagram.branches[branch.name] = branch
+      end
+
+      # Restore commit order
+      diagram.instance_variable_set(:@commit_order, data_hash[:commit_order] || data_hash['commit_order'] || [])
+
+      # Restore current branch name
+      diagram.instance_variable_set(:@current_branch_name,
+                                    data_hash[:current_branch_name] || data_hash['current_branch_name'] || 'master')
+
+      # Recalculate checksum after loading all data
+      diagram.send(:update_checksum!) # Use send to call protected method from class scope
+
+      # Optional: Verify checksum if provided
+      if checksum && diagram.checksum != checksum
+        warn "Checksum mismatch for loaded GitgraphDiagram (version: #{version}). Expected #{checksum}, got #{diagram.checksum}."
+      end
+
+      diagram
+    end
+
+    private
+
+    # Generates a unique ID for a commit if one isn't provided.
+    # Placeholder - could use SHA1/SHA256 of content or simple counter.
+    # Using a simple counter based on commit order for now.
+    def generate_commit_id(parent_ids, message)
+      # Simple approach: use commit count + part of parent hash if available
+      base = @commit_order.size.to_s
+      parent_part = parent_ids.first ? parent_ids.first[0..5] : 'root'
+      # NOTE: This is NOT cryptographically secure or git-like. Just for basic uniqueness.
+      "commit-#{base}-#{parent_part}-#{Digest::SHA1.hexdigest(message || Time.now.to_s)[0..5]}"
+    end
+
+    # Helper to get the head commit ID of the current branch.
+    def current_head_commit_id
+      current_branch = @branches[@current_branch_name]
+      current_branch&.head_commit_id # Returns nil if branch doesn't exist yet
+    end
+
+    # Protected method access for from_h
+    protected :update_checksum!
+  end
+end