Feat: Documenting Dafny Entities (Adding Docstring)

Source/DafnyCore/AST/Function.cs

@@ -7,7 +7,7 @@
 
 namespace Microsoft.Dafny;
 
-public class Function : MemberDecl, TypeParameter.ParentType, ICallable, ICanFormat {
+public class Function : MemberDecl, TypeParameter.ParentType, ICallable, ICanFormat, IHasDocstring {
   public override string WhatKind => "function";
 
   public string GetFunctionDeclarationKeywords(DafnyOptions options) {
@@ -437,4 +437,23 @@ public void Resolve(Resolver resolver) {
 
     resolver.Options.WarnShadowing = warnShadowingOption; // restore the original warnShadowing value
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    if (Body == null) {
+      if (EndToken.TrailingTrivia.Trim() != "") {
+        return EndToken.TrailingTrivia;
+      }
+
+      if (StartToken.LeadingTrivia.Trim() != "") {
+        return StartToken.LeadingTrivia;
+      }
+    } else if (Body.StartToken.Prev.Prev is var tokenWhereTrailingIsDocstring &&
+              tokenWhereTrailingIsDocstring.TrailingTrivia.Trim() != "") {
+      return tokenWhereTrailingIsDocstring.TrailingTrivia;
+    } else if (StartToken is var tokenWhereLeadingIsDocstring &&
+        tokenWhereLeadingIsDocstring.LeadingTrivia.Trim() != "") {
+      return tokenWhereLeadingIsDocstring.LeadingTrivia;
+    }
+    return null;
+  }
 }

Source/DafnyCore/AST/Grammar/TriviaFormatterHelper.cs

@@ -12,11 +12,11 @@ public static bool EndsWithNewline(string s) {
     return EndsWithNewlineRegex.IsMatch(s);
   }
 
-  private static readonly string AnyNewline = @"\r?\n|\r(?!\n)";
+  public static readonly string AnyNewline = @"\r?\n|\r(?!\n)";
 
   private static readonly string NoCommentDelimiter = @"(?:(?!/\*|\*/)[\s\S])*";
 
-  private static readonly string MultilineCommentContent =
+  public static readonly string MultilineCommentContent =
     $@"(?:{NoCommentDelimiter}(?:(?'Open'/\*)|(?'-Open'\*/)))*{NoCommentDelimiter}";
 
   public static readonly Regex NewlineRegex =

Source/DafnyCore/AST/IteratorDecl.cs

@@ -3,7 +3,7 @@
 
 namespace Microsoft.Dafny;
 
-public class IteratorDecl : ClassDecl, IMethodCodeContext {
+public class IteratorDecl : ClassDecl, IMethodCodeContext, IHasDocstring {
   public override string WhatKind { get { return "iterator"; } }
   public readonly List<Formal> Ins;
   public readonly List<Formal> Outs;
@@ -480,4 +480,20 @@ public void Resolve(Resolver resolver) {
       Members.Add(moveNext);
     }
   }
+
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken lastClosingParenthesis = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == ")") {
+        lastClosingParenthesis = token;
+      }
+    }
+
+    if (lastClosingParenthesis != null && lastClosingParenthesis.TrailingTrivia.Trim() != "") {
+      return lastClosingParenthesis.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }

Source/DafnyCore/AST/MemberDecls.cs

@@ -93,7 +93,7 @@ public virtual string FullSanitizedName {
   public virtual IEnumerable<Expression> SubExpressions => Enumerable.Empty<Expression>();
 }
 
-public class Field : MemberDecl, ICanFormat {
+public class Field : MemberDecl, ICanFormat, IHasDocstring {
   public override string WhatKind => "field";
   public readonly bool IsMutable;  // says whether or not the field can ever change values
   public readonly bool IsUserMutable;  // says whether or not code is allowed to assign to the field (IsUserMutable implies IsMutable)
@@ -173,6 +173,14 @@ public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
 
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    if (EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public class SpecialFunction : Function, ICallable {

Source/DafnyCore/AST/Method.cs

@@ -5,7 +5,7 @@
 
 namespace Microsoft.Dafny;
 
-public class Method : MemberDecl, TypeParameter.ParentType, IMethodCodeContext, ICanFormat {
+public class Method : MemberDecl, TypeParameter.ParentType, IMethodCodeContext, ICanFormat, IHasDocstring {
   public override IEnumerable<Node> Children => new Node[] { Body, Decreases }.
     Where(x => x != null).Concat(Ins).Concat(Outs).Concat<Node>(TypeArgs).
     Concat(Req).Concat(Ens).Concat(Mod.Expressions);
@@ -336,4 +336,19 @@ public void Resolve(Resolver resolver) {
       resolver.currentMethod = null;
     }
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken lastClosingParenthesis = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == ")") {
+        lastClosingParenthesis = token;
+      }
+    }
+
+    if (lastClosingParenthesis != null && lastClosingParenthesis.TrailingTrivia.Trim() != "") {
+      return lastClosingParenthesis.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }

Source/DafnyCore/AST/TopLevelDeclarations.cs

@@ -345,7 +345,7 @@ public static Dictionary<TypeParameter, Type> SubstitutionMap(List<TypeParameter
 }
 
 // Represents a submodule declaration at module level scope
-public abstract class ModuleDecl : TopLevelDecl {
+public abstract class ModuleDecl : TopLevelDecl, IHasDocstring {
   public override string WhatKind { get { return "module"; } }
   [FilledInDuringResolution] public ModuleSignature Signature; // filled in topological order.
   public virtual ModuleSignature AccessibleSignature(bool ignoreExports) {
@@ -375,6 +375,27 @@ public override bool IsEssentiallyEmpty() {
     // A module or import is considered "essentially empty" to its parents, but the module is going to be resolved by itself.
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken candidate = null;
+    var tokens = OwnedTokens.Any() ?
+      OwnedTokens :
+      PreResolveChildren.Any() ? PreResolveChildren.First().OwnedTokens : Enumerable.Empty<IToken>();
+    foreach (var token in tokens) {
+      if (token.val == "{") {
+        candidate = token.Prev;
+        if (candidate.TrailingTrivia.Trim() != "") {
+          return candidate.TrailingTrivia;
+        }
+      }
+    }
+
+    if (candidate == null && EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 // Represents module X { ... }
 public class LiteralModuleDecl : ModuleDecl, ICanFormat {
@@ -587,6 +608,14 @@ public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
 
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    if (Tok.TrailingTrivia.Trim() != "") {
+      return Tok.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public class ExportSignature : TokenNode, IHasUsages {
@@ -1349,7 +1378,7 @@ public override IEnumerable<AssumptionDescription> Assumptions() {
   }
 }
 
-public class ClassDecl : TopLevelDeclWithMembers, RevealableTypeDecl, ICanFormat {
+public class ClassDecl : TopLevelDeclWithMembers, RevealableTypeDecl, ICanFormat, IHasDocstring {
   public override string WhatKind { get { return "class"; } }
   public override bool CanBeRevealed() { return true; }
   [FilledInDuringResolution] public bool HasConstructor;  // filled in (early) during resolution; true iff there exists a member that is a Constructor
@@ -1472,6 +1501,24 @@ public virtual bool SetIndent(int indentBefore, TokenNewIndentCollector formatte
 
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken candidate = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == "{") {
+        candidate = token.Prev;
+        if (candidate.TrailingTrivia.Trim() != "") {
+          return candidate.TrailingTrivia;
+        }
+      }
+    }
+
+    if (candidate == null && EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public class DefaultClassDecl : ClassDecl {
@@ -1524,7 +1571,7 @@ public ArrowTypeDecl(List<TypeParameter> tps, Function req, Function reads, Modu
   }
 }
 
-public abstract class DatatypeDecl : TopLevelDeclWithMembers, RevealableTypeDecl, ICallable, ICanFormat {
+public abstract class DatatypeDecl : TopLevelDeclWithMembers, RevealableTypeDecl, ICallable, ICanFormat, IHasDocstring {
   public override bool CanBeRevealed() { return true; }
   public readonly List<DatatypeCtor> Ctors;
 
@@ -1682,6 +1729,16 @@ public bool SetIndent(int indent, TokenNewIndentCollector formatter) {
 
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    foreach (var token in OwnedTokens) {
+      if (token.val == "=" && token.TrailingTrivia.Trim() != "") {
+        return token.TrailingTrivia;
+      }
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public class IndDatatypeDecl : DatatypeDecl {
@@ -1772,7 +1829,7 @@ public Type CreateType(List<Type> typeArgs) {
   }
 }
 
-public class DatatypeCtor : Declaration, TypeParameter.ParentType {
+public class DatatypeCtor : Declaration, TypeParameter.ParentType, IHasDocstring {
   public readonly bool IsGhost;
   public readonly List<Formal> Formals;
   [ContractInvariantMethod]
@@ -1808,6 +1865,14 @@ public string FullName {
       return "#" + EnclosingDatatype.FullName + "." + Name;
     }
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    if (EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 /// <summary>
@@ -1880,6 +1945,7 @@ public CallableWrapper(ICallable callable, bool isGhostContext)
 
   protected ICallable cwInner => (ICallable)inner;
   public IToken Tok => cwInner.Tok;
+
   public string WhatKind => cwInner.WhatKind;
   public string NameRelativeToModule => cwInner.NameRelativeToModule;
   public Specification<Expression> Decreases => cwInner.Decreases;
@@ -1903,6 +1969,10 @@ public class DontUseICallable : ICallable {
   public string FullSanitizedName { get { throw new cce.UnreachableException(); } }
   public bool AllowsNontermination { get { throw new cce.UnreachableException(); } }
   public IToken Tok { get { throw new cce.UnreachableException(); } }
+  public string GetDocstring(DafnyOptions options) {
+    throw new cce.UnreachableException();
+  }
+
   public string NameRelativeToModule { get { throw new cce.UnreachableException(); } }
   public Specification<Expression> Decreases { get { throw new cce.UnreachableException(); } }
   public bool InferredDecreases {
@@ -1939,7 +2009,7 @@ public NoContext(ModuleDefinition module) {
   public bool AllowsAllocation => true;
 }
 
-public class OpaqueTypeDecl : TopLevelDeclWithMembers, TypeParameter.ParentType, RevealableTypeDecl, ICanFormat {
+public class OpaqueTypeDecl : TopLevelDeclWithMembers, TypeParameter.ParentType, RevealableTypeDecl, ICanFormat, IHasDocstring {
   public override string WhatKind { get { return "opaque type"; } }
   public override bool CanBeRevealed() { return true; }
   public readonly TypeParameter.TypeParameterCharacteristics Characteristics;
@@ -2020,6 +2090,25 @@ public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
 
     return true;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken openingBlock = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == "{") {
+        openingBlock = token;
+      }
+    }
+
+    if (openingBlock != null && openingBlock.Prev.TrailingTrivia.Trim() != "") {
+      return openingBlock.Prev.TrailingTrivia;
+    }
+
+    if (openingBlock == null && EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public interface RedirectingTypeDecl : ICallable {
@@ -2091,7 +2180,7 @@ public interface RevealableTypeDecl {
   TypeDeclSynonymInfo SynonymInfo { get; set; }
 }
 
-public class NewtypeDecl : TopLevelDeclWithMembers, RevealableTypeDecl, RedirectingTypeDecl {
+public class NewtypeDecl : TopLevelDeclWithMembers, RevealableTypeDecl, RedirectingTypeDecl, IHasDocstring {
   public override string WhatKind { get { return "newtype"; } }
   public override bool CanBeRevealed() { return true; }
   public Type BaseType { get; set; } // null when refining
@@ -2177,9 +2266,27 @@ public override bool IsEssentiallyEmpty() {
     // A "newtype" is not considered "essentially empty", because it always has a parent type to be resolved.
     return false;
   }
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken candidate = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == "{") {
+        candidate = token.Prev;
+        if (candidate.TrailingTrivia.Trim() != "") {
+          return candidate.TrailingTrivia;
+        }
+      }
+    }
+
+    if (candidate == null && EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
-public abstract class TypeSynonymDeclBase : TopLevelDecl, RedirectingTypeDecl {
+public abstract class TypeSynonymDeclBase : TopLevelDecl, RedirectingTypeDecl, IHasDocstring {
   public override string WhatKind { get { return "type synonym"; } }
   public TypeParameter.TypeParameterCharacteristics Characteristics;  // the resolver may change the .EqualitySupport component of this value from Unspecified to InferredRequired (for some signatures that may immediately imply that equality support is required)
   public bool SupportsEquality {
@@ -2273,6 +2380,27 @@ public override bool IsEssentiallyEmpty() {
     // A synonym/subset type is not considered "essentially empty", because it always has a parent type to be resolved.
     return false;
   }
+
+
+
+  protected override string GetTriviaContainingDocstring() {
+    IToken openingBlock = null;
+    foreach (var token in OwnedTokens) {
+      if (token.val == "{") {
+        openingBlock = token;
+      }
+    }
+
+    if (openingBlock != null && openingBlock.Prev.TrailingTrivia.Trim() != "") {
+      return openingBlock.Prev.TrailingTrivia;
+    }
+
+    if (openingBlock == null && EndToken.TrailingTrivia.Trim() != "") {
+      return EndToken.TrailingTrivia;
+    }
+
+    return GetTriviaContainingDocstringFromStartTokenOrNull();
+  }
 }
 
 public class TypeSynonymDecl : TypeSynonymDeclBase, RevealableTypeDecl {