From 6ce42823cb4fa397b60fda53625009420a7f6bd2 Mon Sep 17 00:00:00 2001
From: topecongiro <seuchida@gmail.com>
Date: Fri, 10 Nov 2017 17:08:57 +0900
Subject: [PATCH] Add opt_snippet() and opt_span_after()

---
 src/codemap.rs | 8 ++++++++
 src/visitor.rs | 4 ++++
 2 files changed, 12 insertions(+)

diff --git a/src/codemap.rs b/src/codemap.rs
index f62bb16dfa2..efbe50bcf64 100644
--- a/src/codemap.rs
+++ b/src/codemap.rs
@@ -34,6 +34,7 @@ pub trait SpanUtils {
     fn span_after(&self, original: Span, needle: &str) -> BytePos;
     fn span_after_last(&self, original: Span, needle: &str) -> BytePos;
     fn span_before(&self, original: Span, needle: &str) -> BytePos;
+    fn opt_span_after(&self, original: Span, needle: &str) -> Option<BytePos>;
 }
 
 pub trait LineRangeUtils {
@@ -70,6 +71,13 @@ impl SpanUtils for CodeMap {
 
         original.lo() + BytePos(offset as u32)
     }
+
+    fn opt_span_after(&self, original: Span, needle: &str) -> Option<BytePos> {
+        let snippet = self.span_to_snippet(original).ok()?;
+        let offset = snippet.find_uncommented(needle)? + needle.len();
+
+        Some(original.lo() + BytePos(offset as u32))
+    }
 }
 
 impl LineRangeUtils for CodeMap {
diff --git a/src/visitor.rs b/src/visitor.rs
index c613b3b8a55..8e8cd379e82 100644
--- a/src/visitor.rs
+++ b/src/visitor.rs
@@ -510,6 +510,10 @@ impl<'a> FmtVisitor<'a> {
         }
     }
 
+    pub fn opt_snippet(&self, span: Span) -> Option<String> {
+        self.codemap.span_to_snippet(span).ok()
+    }
+
     pub fn snippet(&self, span: Span) -> String {
         match self.codemap.span_to_snippet(span) {
             Ok(s) => s,