R/if_switch_linter.R

In lintr: A 'Linter' for R Code

#' Require usage of switch() over repeated if/else blocks
#'
#' [switch()] statements in R are used to delegate behavior based
#'   on the value of some input scalar string, e.g.
#'   `switch(x, a = 1, b = 3, c = 7, d = 8)` will be one of
#'   `1`, `3`, `7`, or `8`, depending on the value of `x`.
#'
#' This can also be accomplished by repeated `if`/`else` statements like
#'   so: `if (x == "a") 1 else if (x == "b") 2 else if (x == "c") 7 else 8`
#'   (implicitly, the last `else` assumes x only takes 4 possible values),
#'   but this is more cluttered and slower (note that `switch()` takes the same
#'   time to evaluate regardless of the value of `x`, and is faster even
#'   when `x` takes the first value (here `a`), and that the `if`/`else`
#'   approach is roughly linear in the number of conditions that need to
#'   be evaluated, here up to 3 times).
#'
#' @param max_branch_lines,max_branch_expressions Integer, default 0 indicates "no maximum".
#'   If set any `if`/`else if`/.../`else` chain where any branch occupies more than
#'   this number of lines (resp. expressions) will not be linted. The conjugate
#'   applies to `switch()` statements -- if these parameters are set, any `switch()`
#'   statement with any overly-complicated branches will be linted. See examples.
#'
#' @examples
#' # will produce lints
#' lint(
#'   text = "if (x == 'a') 1 else if (x == 'b') 2 else 3",
#'   linters = if_switch_linter()
#' )
#'
#' code <- paste(
#'   "if (x == 'a') {",
#'   "  1",
#'   "} else if (x == 'b') {",
#'   "  2",
#'   "} else if (x == 'c') {",
#'   "  y <- x",
#'   "  z <- sqrt(match(y, letters))",
#'   "  z",
#'   "}",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter()
#' )
#'
#' code <- paste(
#'   "if (x == 'a') {",
#'   "  1",
#'   "} else if (x == 'b') {",
#'   "  2",
#'   "} else if (x == 'c') {",
#'   "  y <- x",
#'   "  z <- sqrt(",
#'   "    match(y, letters)",
#'   "  )",
#'   "  z",
#'   "}",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter()
#' )
#'
#' code <- paste(
#'   "switch(x,",
#'   "  a = {",
#'   "    1",
#'   "    2",
#'   "    3",
#'   "  },",
#'   "  b = {",
#'   "    1",
#'   "    2",
#'   "  }",
#'   ")",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter(max_branch_lines = 2L)
#' )
#'
#' # okay
#' lint(
#'   text = "switch(x, a = 1, b = 2, 3)",
#'   linters = if_switch_linter()
#' )
#'
#' # switch() version not as clear
#' lint(
#'   text = "if (x == 'a') 1 else if (x == 'b' & y == 2) 2 else 3",
#'   linters = if_switch_linter()
#' )
#'
#' code <- paste(
#'   "if (x == 'a') {",
#'   "  1",
#'   "} else if (x == 'b') {",
#'   "  2",
#'   "} else if (x == 'c') {",
#'   "  y <- x",
#'   "  z <- sqrt(match(y, letters))",
#'   "  z",
#'   "}",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter(max_branch_lines = 2L)
#' )
#'
#' code <- paste(
#'   "if (x == 'a') {",
#'   "  1",
#'   "} else if (x == 'b') {",
#'   "  2",
#'   "} else if (x == 'c') {",
#'   "  y <- x",
#'   "  z <- sqrt(",
#'   "    match(y, letters)",
#'   "  )",
#'   "  z",
#'   "}",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter(max_branch_expressions = 2L)
#' )
#'
#' code <- paste(
#'   "switch(x,",
#'   "  a = {",
#'   "    1",
#'   "    2",
#'   "    3",
#'   "  },",
#'   "  b = {",
#'   "    1",
#'   "    2",
#'   "  }",
#'   ")",
#'   sep = "\n"
#' )
#' writeLines(code)
#' lint(
#'   text = code,
#'   linters = if_switch_linter(max_branch_lines = 3L)
#' )
#'
#' @evalRd rd_tags("if_switch_linter")
#' @seealso [linters] for a complete list of linters available in lintr.
#' @export
if_switch_linter <- function(max_branch_lines = 0L, max_branch_expressions = 0L) {
  equal_str_cond <- "expr[1][EQ and expr/STR_CONST[string-length(text()) > 2]]"

  if (max_branch_lines > 0L || max_branch_expressions > 0L) {
    complexity_cond <- xp_or(c(
      if (max_branch_lines > 0L) {
        c(
          paste("OP-LEFT-BRACE and (OP-RIGHT-BRACE/@line2 - OP-LEFT-BRACE/@line1 > 1 +", max_branch_lines, ")"),
          paste("not(OP-LEFT-BRACE) and (@line2 - @line1 >=", max_branch_lines, ")")
        )
      },
      if (max_branch_expressions > 0L) paste("OP-LEFT-BRACE and count(expr) >", max_branch_expressions)
    ))
    branch_expr_cond <- xp_and(c(
      xp_or(
        # if (x) { <this expr> } ...
        xp_and("preceding-sibling::IF", "position() = 2"),
        # if (x) { ... } else { <this expr> }
        xp_and("preceding-sibling::ELSE", "not(IF)")
      ),
      complexity_cond
    ))
    max_lines_cond <- glue(".//expr[{branch_expr_cond}]")

    switch_xpath <- glue("
    parent::expr[expr[
      position() > 2
      and {complexity_cond}
    ]]")
  } else {
    max_lines_cond <- "false"

    switch_xpath <- NULL
  }

  if_xpath <- glue("
  //IF
    /parent::expr[
      not(preceding-sibling::IF)
      and {equal_str_cond}
      and ELSE/following-sibling::expr[
        IF
        and {equal_str_cond}
        and ELSE/following-sibling::expr[IF and {equal_str_cond}]
      ]
      and not({ max_lines_cond })
    ]
  ")

  Linter(linter_level = "expression", function(source_expression) {
    xml <- source_expression$xml_parsed_content

    bad_expr <- xml_find_all_(xml, if_xpath)

    bad_expr_clean <- strip_comments_from_subtree(bad_expr)
    expr_all_equal <- vapply(bad_expr_clean, if_else_chain_expr_is_unique, logical(1L))
    bad_expr <- bad_expr[expr_all_equal]

    # Exclude empty strings, which can't be used as switch() case names
    nonempty <- vapply(bad_expr, function(expr) {
      str_nodes <- if_else_chain_strings(expr)
      all(vapply(str_nodes, function(n) nzchar(get_r_string(n)), logical(1L)))
    }, logical(1L))
    bad_expr <- bad_expr[nonempty]

    lints <- xml_nodes_to_lints(
      bad_expr,
      source_expression = source_expression,
      lint_message = paste(
        "Prefer switch() statements over repeated if/else equality tests,",
        "e.g., switch(x, a = 1, b = 2) over",
        'if (x == "a") 1 else if (x == "b") 2.'
      ),
      type = "warning"
    )

    if (!is.null(switch_xpath)) {
      xml_calls <- source_expression$xml_find_function_calls("switch")
      switch_expr <- xml_find_all_(xml_calls, switch_xpath)

      lints <- c(lints, xml_nodes_to_lints(
        switch_expr,
        source_expression = source_expression,
        lint_message = "Prefer repeated if/else statements over overly-complicated switch() statements.",
        type = "warning"
      ))
    }

    lints
  })
}

# Extract STR_CONST nodes from equality conditions in an if/else if chain
if_else_chain_strings <- function(expr) {
  str_nodes <- list()
  first <- xml_find_first_(expr, "IF/following-sibling::expr[1][EQ]/expr/STR_CONST")
  if (!is.na(first)) str_nodes <- c(str_nodes, list(first))
  current <- expr
  repeat {
    else_if <- xml_find_first_(current, "ELSE/following-sibling::expr[IF]")
    if (is.na(else_if)) break
    current <- else_if
    str_node <- xml_find_first_(current, "IF/following-sibling::expr[1][EQ]/expr/STR_CONST")
    if (!is.na(str_node)) str_nodes <- c(str_nodes, list(str_node))
  }
  str_nodes
}

# Check that equality conditions in an if/else if chain use the same expression
if_else_chain_expr_is_unique <- function(expr) {
  expr_nodes <- character()
  first <- xml_find_first_(expr, "IF/following-sibling::expr[1][EQ]/expr[not(STR_CONST)]")
  if (!is.na(first)) expr_nodes <- c(expr_nodes, xml_text(first))
  current <- expr
  repeat {
    else_if <- xml_find_first_(current, "ELSE/following-sibling::expr[IF]")
    if (is.na(else_if)) break
    current <- else_if
    expr_node <- xml_find_first_(current, "IF/following-sibling::expr[1][EQ]/expr[not(STR_CONST)]")
    if (!is.na(expr_node)) expr_nodes <- c(expr_nodes, xml_text(expr_node))
  }
  length(unique(expr_nodes)) == 1L
}