Convert stringmask macro to the scalac compiler plugin

annotation/src/main/scala/com/softwaremill/stringmask/annotation/mask.scala

@@ -0,0 +1,5 @@
+package com.softwaremill.stringmask.annotation
+
+import scala.annotation.StaticAnnotation
+
+class mask extends StaticAnnotation

build.sbt

@@ -1,42 +1,75 @@
-organization := "com.softwaremill.stringmask"
-name := "stringmask"
+import sbt.Keys._
 
-version := "1.1.0-SNAPSHOT"
+import scalariform.formatter.preferences._
+
+lazy val commonSettings = scalariformSettings ++ Seq(
+    organization := "com.softwaremill.stringmask",
+    version := "2.0.0-SNAPSHOT",
+    crossScalaVersions := Seq("2.10.6", "2.11.8"),
+    scalaVersion := "2.11.8",
+    ScalariformKeys.preferences in ThisBuild := ScalariformKeys.preferences.value
+        .setPreference(DoubleIndentClassDeclaration, true)
+        .setPreference(PreserveSpaceBeforeArguments, true)
+        .setPreference(CompactControlReadability, true)
+        .setPreference(SpacesAroundMultiImports, false)
+)
+
+lazy val root = (project in file("."))
+    .aggregate(annotation, scalacPlugin, tests)
+    .settings(commonSettings)
+    .settings(name := "scalac-stringmask-plugin")
 
-crossScalaVersions := Seq("2.10.6", "2.11.8")
 
-scalaVersion := "2.11.8"
+lazy val annotation = (project in file("annotation"))
+    .settings(commonSettings)
+    .settings(name := "stringmask-annotation")
+    .settings(publishArtifact := true)
 
-libraryDependencies ++= Seq(
-  "org.scala-lang" % "scala-reflect" % scalaVersion.value,
-  "org.scala-lang" % "scala-compiler" % scalaVersion.value,
-  "org.typelevel" %% "macro-compat" % "1.1.1",
-  "org.scalatest" %% "scalatest" % "2.2.6" % "test"
+lazy val scalacPlugin = (project in file("scalacPlugin"))
+    .dependsOn(annotation)
+    .settings(commonSettings)
+    .settings(
+        name := "stringmask-scalac-plugin",
+        exportJars := true
+    )
+    .settings(
+        libraryDependencies ++= Seq(
+            "org.scala-lang" % "scala-compiler" % scalaVersion.value
+        )
+    )
+
+lazy val tests = (project in file("tests"))
+    .dependsOn(scalacPlugin)
+    .settings(commonSettings)
+    .settings(
+        scalacOptions <+= (artifactPath in(scalacPlugin, Compile, packageBin)).map { file =>
+            s"-Xplugin:${file.getAbsolutePath}"
+        }
+    ).settings(
+    libraryDependencies ++= Seq(
+        "org.scalatest" %% "scalatest" % "2.2.6"
+    )
 )
 
-publishTo := {
-  val nexus = "https://oss.sonatype.org/"
-  if (version.value.trim.endsWith("SNAPSHOT"))
-    Some("snapshots" at nexus + "content/repositories/snapshots")
-  else
-    Some("releases" at nexus + "service/local/staging/deploy/maven2")
-}
-credentials += Credentials(Path.userHome / ".ivy2" / ".credentials")
-publishMavenStyle := true
-publishArtifact in Test := false
+publishMavenStyle in ThisBuild := true
+publishArtifact in ThisBuild := true
+publishArtifact in Test in ThisBuild := false
 pomIncludeRepository := { _ => false }
-pomExtra :=
-  <scm>
-    <url>git@github.com:softwaremill/stringmask.git</url>
-    <connection>scm:git:git@github.com:softwaremill/stringmask.git</connection>
-  </scm>
-    <developers>
-      <developer>
-        <id>kciesielski</id>
-        <name>Krzysztof Ciesielski</name>
-      </developer>
-    </developers>
-licenses := ("Apache2", new java.net.URL("http://www.apache.org/licenses/LICENSE-2.0.txt")) :: Nil
-homepage := Some(new java.net.URL("http://www.softwaremill.com"))
-
-addCompilerPlugin("org.scalamacros" % "paradise" % "2.1.0" cross CrossVersion.full)
+pomExtra in ThisBuild :=
+    <scm>
+        <url>git@github.com:softwaremill/stringmask.git</url>
+        <connection>scm:git:git@github.com:softwaremill/stringmask.git</connection>
+    </scm>
+        <developers>
+            <developer>
+                <id>kciesielski</id>
+                <name>Krzysztof Ciesielski</name>
+            </developer>
+            <developer>
+                <id>mkubala</id>
+                <name>Marcin Kubala</name>
+            </developer>
+        </developers>
+
+licenses in ThisBuild := ("Apache2", new java.net.URL("http://www.apache.org/licenses/LICENSE-2.0.txt")) :: Nil
+homepage in ThisBuild := Some(new java.net.URL("http://www.softwaremill.com"))

project/plugins.sbt

@@ -0,0 +1 @@
+addSbtPlugin("org.scalariform" % "sbt-scalariform" % "1.4.0")

scalacPlugin/src/main/resources/scalac-plugin.xml

@@ -0,0 +1,4 @@
+<plugin>
+	<name>stringmask</name>
+	<classname>com.softwaremill.stringmask.StringMaskPlugin</classname>
+</plugin>

scalacPlugin/src/main/scala/com/softwaremill/stringmask/StringMaskPlugin.scala

@@ -0,0 +1,13 @@
+package com.softwaremill.stringmask
+
+import com.softwaremill.stringmask.components.StringMaskComponent
+
+import scala.tools.nsc.Global
+import scala.tools.nsc.plugins.{ Plugin, PluginComponent }
+
+class StringMaskPlugin(val global: Global) extends Plugin {
+
+  override val name: String = "stringmask"
+  override val description: String = "StringMask compiler plugin"
+  override val components: List[PluginComponent] = List(new StringMaskComponent(global))
+}

scalacPlugin/src/main/scala/com/softwaremill/stringmask/components/StringMaskComponent.scala

@@ -0,0 +1,83 @@
+package com.softwaremill.stringmask.components
+
+import scala.tools.nsc.Global
+import scala.tools.nsc.plugins.PluginComponent
+import scala.tools.nsc.transform.Transform
+
+class StringMaskComponent(val global: Global) extends PluginComponent with Transform {
+
+  override val phaseName: String = "stringmask"
+
+  override val runsAfter: List[String] = List("parser")
+
+  // When runs after typer phase, might have problems with accessing ValDefs' annotations.
+  override val runsRightAfter: Option[String] = Some("parser")
+
+  override protected def newTransformer(unit: global.CompilationUnit): global.Transformer = ToStringMaskerTransformer
+
+  import global._
+
+  object ToStringMaskerTransformer extends Transformer {
+
+    override def transform(tree: global.Tree): global.Tree = {
+      val transformedTree = super.transform(tree)
+      transformedTree match {
+        case classDef: ClassDef if isAnnotatedCaseClass(classDef) =>
+          extractParamsAnnotatedWithMask(classDef)
+            .map(buildNewToStringTree(classDef.name))
+            .map(overrideToStringDef(classDef))
+            .getOrElse(transformedTree)
+        case oth => transformedTree
+      }
+    }
+
+    private def isAnnotatedCaseClass(classDef: ClassDef): Boolean =
+      classDef.mods.isCase && !containsCustomToStringDef(classDef)
+
+    private def containsCustomToStringDef(classDef: global.ClassDef): Boolean =
+      classDef.impl.body.exists {
+        case d: DefDef => d.name.decode == "toString"
+        case _ => false
+      }
+
+    private def extractParamsAnnotatedWithMask(classDef: ClassDef): Option[List[Tree]] =
+      classDef.impl.body.collectFirst {
+        case d: DefDef if d.name.decode == "<init>" && d.vparamss.headOption.exists(containsMaskedParams) =>
+          d.vparamss.headOption.map { firstParamsGroup =>
+            firstParamsGroup.foldLeft(List.empty[Tree]) {
+              case (accList, fieldTree) =>
+                val newFieldTree = if (hasMaskAnnotation(fieldTree)) {
+                  Literal(Constant("***"))
+                } else {
+                  Apply(Select(Ident(fieldTree.name), "toString"), Nil)
+                }
+                accList :+ newFieldTree
+            }
+          }.getOrElse(Nil)
+      }
+
+    private def containsMaskedParams(params: List[ValDef]): Boolean =
+      params.exists(hasMaskAnnotation)
+
+    private def hasMaskAnnotation(param: ValDef): Boolean =
+      param.mods.hasAnnotationNamed("mask")
+
+    private def overrideToStringDef(classDef: global.ClassDef)(newToStringImpl: Tree): Tree = {
+      val className = classDef.name
+      global.inform(s"overriding $className.toString")
+      val newBody = newToStringImpl :: classDef.impl.body
+      val newImpl = Template(classDef.impl.parents, classDef.impl.self, newBody)
+      ClassDef(classDef.mods, className, classDef.tparams, newImpl)
+    }
+
+    private def buildNewToStringTree(className: TypeName)(fields: List[Tree]): Tree = {
+      val treesAsTuple = Apply(Select(Ident("scala"), "Tuple" + fields.length), fields)
+      val typeNameStrTree = Literal(Constant(className.toString))
+
+      DefDef(Modifiers(Flag.OVERRIDE), "toString": TermName, List(), List(), TypeTree(),
+        Apply(Select(typeNameStrTree, "$plus": TermName), List(treesAsTuple)))
+    }
+
+  }
+
+}

tests/src/test/scala/com/softwaremill/stringmask/testing/BasicCasesSpec.scala

@@ -0,0 +1,51 @@
+package com.softwaremill.stringmask.testing
+
+import com.softwaremill.stringmask.annotation.mask
+import org.scalatest.{ Matchers, WordSpec }
+
+class BasicCasesSpec extends WordSpec with Matchers {
+  "StringMask" should {
+    "not override toString" when {
+      "none of the class parameters is annotated with @mask" in {
+        case class Planet(name: String)
+
+        Planet("earth").toString should equal("Planet(earth)")
+      }
+
+      "@mask annotation is applied, but not in the first parameters list" in {
+        case class TrickyCurrying(firstArg: String)(@mask secretKey: String)
+
+        TrickyCurrying("red")("tomato").toString should equal("TrickyCurrying(red)")
+      }
+
+      "class is implementing a custom toString method" in {
+        case class UserWithCustomToString(name: String, @mask password: String) {
+          override def toString: String = s"I love potatoes"
+        }
+
+        UserWithCustomToString("James", "secretPass").toString should equal("I love potatoes")
+      }
+
+      "applied to regular class" in {
+        new TestClasses.FavouriteMug("white", 0.4f, "yerba mate").toString should startWith("com.softwaremill.stringmask.testing.TestClasses$FavouriteMug@")
+      }
+
+    }
+
+    "mask confidential fields" when {
+
+      "applied to case classes" in {
+        case class CasualUser(name: String, @mask password: String)
+
+        CasualUser("James", "secretPass").toString should equal("CasualUser(James,***)")
+      }
+
+    }
+  }
+}
+
+object TestClasses {
+
+  class FavouriteMug(color: String, volume: Float, @mask content: String)
+
+}