1 files changed, 15 insertions, 11 deletions

diff --git a/src/emitter.rs b/src/emitter.rs
index 9fdf967..db3da78 100644
--- a/src/emitter.rs
+++ b/src/emitter.rs
@@ -55,7 +55,7 @@ pub trait Emitter<O> {
     fn init_end_tag(&mut self, tag_offset: O, name_offset: O);
 
     /// Set the _current token_ to a comment.
-    fn init_comment(&mut self, data_offset: O);
+    fn init_comment(&mut self, data_start_offset: O);
 
     /// Emit the _current token_, assuming it is a tag.
     ///
@@ -71,7 +71,7 @@ pub trait Emitter<O> {
     /// Emit the _current token_, assuming it is a comment.
     ///
     /// If the current token is not a comment, this method may panic.
-    fn emit_current_comment(&mut self, offset: O);
+    fn emit_current_comment(&mut self, data_end_offset: O);
 
     /// Emit the _current token_, assuming it is a doctype.
     ///
@@ -309,10 +309,10 @@ impl<O: Offset> Emitter<O> for DefaultEmitter<O> {
         self.seen_attributes.clear();
     }
 
-    fn init_comment(&mut self, data_offset: O) {
+    fn init_comment(&mut self, data_start_offset: O) {
         self.current_token = Some(Token::Comment(Comment {
             data: String::new(),
-            data_offset,
+            data_span: data_start_offset..O::default(),
         }));
     }
     fn emit_current_tag(&mut self, offset: O) {
@@ -334,10 +334,14 @@ impl<O: Offset> Emitter<O> for DefaultEmitter<O> {
         }
         self.emit_token(token);
     }
-    fn emit_current_comment(&mut self, _offset: O) {
-        let comment = self.current_token.take().unwrap();
-        debug_assert!(matches!(comment, Token::Comment(_)));
-        self.emit_token(comment);
+    fn emit_current_comment(&mut self, data_end_offset: O) {
+        let mut token = self.current_token.take().unwrap();
+        if let Token::Comment(comment) = &mut token {
+            comment.data_span.end = data_end_offset;
+        } else {
+            debug_assert!(false);
+        }
+        self.emit_token(token);
     }
 
     fn emit_current_doctype(&mut self, offset: O) {
@@ -572,13 +576,13 @@ pub struct Comment<O> {
     /// The text within the comment.
     pub data: String,
     /// The source offset of the comment data.
-    pub data_offset: O,
+    pub data_span: Range<O>,
 }
 
 impl<O: Offset> Comment<O> {
-    /// Calculates the span for the comment data and returns it.
+    /// Returns the span for the comment data.
     pub fn data_span(&self) -> Range<O> {
-        self.data_offset..self.data_offset + self.data.len()
+        self.data_span.clone()
     }
 }