<commit_before>-- Imports / module wrapper
local Pkg = {}
local std = _G
local error = error
local print = print
local type = type
local tostring = tostring
local getmetatable = getmetatable
local setmetatable = setmetatable
local string = string
local util = require("util")
local Set = require("set").Set
local parser = require("parser")
local color = require("color")
local build = require("build")
setfenv(1, Pkg)

local ENTITY_FIELD = parser.ENTITY_FIELD
local TAG_FIELD = "tag"
local EAV_TAG = "eav"

-- Utilities
local nothing = {}

function formatQueryNode(node, indent)
   indent = indent or 0
   local padding = string.rep("  ", indent)
   local result = padding .. node.type
   if node.type == "query" then
      result = result .. "<" .. (node.name or "unnamed") .. ">"
      if node.unpacked then
         result = result .. "{\n"
         for ix, guy in std.ipairs(node.unpacked) do
            result = result .. padding .. "  " .. ix .. ". " .. tostring(guy) .. ",\n"
         end
         result = result .. padding .. "}"
      elseif node.dependencyGraph then
         result = result .. tostring(node.dependencyGraph)
      end
   elseif node.type == "constant" then
     result = result .. "<" .. node.constant .. ">"
   elseif node.type == "variable" then
      result = result .. "<" .. (node.name or "unnamed") .. ">"
   elseif node.type == "binding" then
      result = result .. "{" .. tostring(node.field) .. " -> "
      if node.constant then
         result = result .. tostring(node.constant.constant)
      elseif node.variable then
         result = result .. formatQueryNode(node.variable)
      end
      return result .. "}"
   elseif node.type == "object" then
      result = result .. "{"
      for _, binding in std.ipairs(node.bindings) do
         result = result .. formatQueryNode(binding) .. ", "
      end
      return result .. "}"
   elseif node.type == "mutate" then
      result = result .. "<" .. node.operator .. ">{"
      for _, binding in std.ipairs(node.bindings) do
         result = result .. formatQueryNode(binding) .. ", "
      end
      return result .. "}"
   elseif node.type == "union" or node.type == "choose" then
      result = result .. "{\n"
      for _, query in std.ipairs(node.queries) do
         result = result .. formatQueryNode(query, 1) .. ",\n"
      end
      return result .. "}"
   elseif node.type == "expression" then
      result = result .. " " .. node.operator .. "("
      for _, binding in std.ipairs(node.bindings) do
         result = result .. binding.field .. " = " .. formatQueryNode(binding.variable or binding.constant) .. ", "
      end
      result = string.sub(result, 1, -3) .. ")"
   end
   return result
end

local DefaultNodeMeta = {}
DefaultNodeMeta.__tostring = formatQueryNode

local function applyDefaultMeta(_, node)
   if type(node) == "table" and getmetatable(node) == nil then
      setmetatable(node, DefaultNodeMeta)
   end
end

-- Dependency Graph

DependencyGraph = {}

function DependencyGraph:new(obj)
   obj = obj or {}
   -- Essential state
   obj.unsorted = obj.unsorted or Set:new() -- set of nodes that need to be ordered
   obj.sorted = obj.sorted or {} -- append-only set of ordered nodes

   obj.unsatisfied = obj.unsatisfied or {} -- Number of terms required but not bound per node
   obj.dependents = obj.dependents or {} -- Sets of nodes depending on a term
   obj.strongDependents = obj.strongDependents or {} -- Sets of nodes strongly depending on (requiring stabilization of) a term
   obj.bound = obj.bound or Set:new() -- Set of terms bound by the currently ordered set of nodes
   obj.terms = obj.terms or Set:new() -- Set of all terms provided by any node in the graph
   obj.termGroups = obj.termGroups or Set:new() -- Set of sets of terms that are in some way joined

   setmetatable(obj, self)
   self.__index = self
   return obj
end

-- Get the variables this dgraph depends on for reification
function DependencyGraph:depends()
   local depends = Set:new()
   for term in std.pairs(self.dependents) do
      depends:add(term)
   end
   return depends / self.bound
end

function DependencyGraph:provided()
   return self.bound
end

function DependencyGraph:isSorted()
   return self.unsorted:length() == 0
end

function DependencyGraph:addObjectNode(node)
   local produces = Set:new()
   local depends = Set:new()
   for _, binding in std.ipairs(node.bindings or nothing) do
      if binding.variable then
         produces:add(binding.variable)
         depends:add(binding.variable)
      end
   end
   return self:add(node, depends, produces)
end

function DependencyGraph:addMutateNode(node)
   local produces = Set:new()
   local depends = Set:new()
   local weakDepends = Set:new()
   for _, binding in std.ipairs(node.bindings or nothing) do
      -- If the entity term isn't bound, the mutation produces it
      if binding.field == ENTITY_FIELD and not self.terms[binding.variable] then
         produces:add(binding.variable)
      end

      -- If the binding is bound on a variable that is produced in the query, it becomes a dependency of the query.
      if binding.variable then
         weakDepends:add(binding.variable)
      end
   end

   return self:add(node, depends, produces, weakDepends)
end

function DependencyGraph:addExpressionNode(node)
   -- TODO also need to consider projections and groupings as dependencies
   -- TODO handle productions other than just "return", this will require schemas
   local produces = Set:new()
   local depends = Set:new()
   local maybeDepends
   local strongDepend
   local anyDepends
   for _, binding in std.ipairs(node.bindings) do
     if binding.field == "return" and binding.variable then
       produces:add(binding.variable)
     elseif binding.variable then
       depends:add(binding.variable)
     end
   end
   if node.operator == "=" and produces:length() == 0 then
      anyDepends = depends
      produces = depends:clone()
      depends = Set:new()
   end
   return self:add(node, depends, produces, maybeDepends, strongDepends, anyDepends)
end

function DependencyGraph:addSubqueryNode(node)
   local provides = Set:new()
   local depends = Set:new()

   for _, var in std.pairs(node.outputs) do
     provides:add(var)
   end

   for _, body in std.ipairs(node.queries) do
      local subgraph = DependencyGraph:fromQueryGraph(body, self.terms:clone(), self.terms:clone())
      provides:union(subgraph:provided(), true)
      depends:union(subgraph:depends(), true)
   end
   return self:add(node, depends, provides)
end

function DependencyGraph:fromQueryGraph(query, terms, bound)
   local uniqueCounter = 0
   local dgraph = self
   if getmetatable(dgraph) ~= DependencyGraph then
      -- @FIXME: this naive injection strategy lets objects be scheduled willy-nilly after mutates...
      -- Not tracking this information is currently incidentally correct, but cannot be relied upon to order subqueries properly within their parents
      -- dgraph = self:new{terms = terms, bound = bound}
      dgraph = self:new()

   end
   dgraph.query = query
   query.dependencyGraph = dgraph

   util.walk(query, applyDefaultMeta)

   for _, node in std.ipairs(query.expressions or nothing) do
      dgraph:addExpressionNode(node)
   end

   for _, node in std.ipairs(query.nots or nothing) do
      dgraph:addSubqueryNode(node)
   end

   for _, node in std.ipairs(query.objects or nothing) do
      dgraph:addObjectNode(node)
   end

   for _, node in std.ipairs(query.unions or nothing) do
      dgraph:addSubqueryNode(node)
   end

   for _, node in std.ipairs(query.chooses or nothing) do
      dgraph:addSubqueryNode(node)
   end

   -- If the ENTITY variable of this mutate isn't already bound and it's a child of another mutate,
   -- then we need to uniquify it to prevent unification with other mutates.
   for _, node in std.ipairs(query.mutates or nothing) do
      if not dgraph.terms[node.variable] then
         if node.parent and node.parent.type == "mutate" then
            local old = node.variable
            node.variable = util.shallowCopy(node.variable)
            node.variable.name = "$$tmp" .. "-" .. node.variable.name .. "-" .. uniqueCounter
            uniqueCounter = uniqueCounter + 1
            for _, binding in std.pairs(node.bindings) do
               if binding.type == "binding" and binding.variable == old then
                  binding.variable = node.variable
               end
            end
            local parent = node.parent
            for _, binding in std.pairs(parent.bindings) do
               if binding.type == "binding" and binding.variable == old then
                  binding.variable = node.variable
               end
            end
         end
      end
   end<commit_msg>fix assignment binding constants and union/choose dependencies<commit_after>-- Imports / module wrapper
local Pkg = {}
local std = _G
local error = error
local print = print
local type = type
local tostring = tostring
local getmetatable = getmetatable
local setmetatable = setmetatable
local string = string
local util = require("util")
local Set = require("set").Set
local parser = require("parser")
local color = require("color")
local build = require("build")
setfenv(1, Pkg)

local ENTITY_FIELD = parser.ENTITY_FIELD
local TAG_FIELD = "tag"
local EAV_TAG = "eav"

-- Utilities
local nothing = {}

function formatQueryNode(node, indent)
   indent = indent or 0
   local padding = string.rep("  ", indent)
   local result = padding .. node.type
   if node.type == "query" then
      result = result .. "<" .. (node.name or "unnamed") .. ">"
      if node.unpacked then
         result = result .. "{\n"
         for ix, guy in std.ipairs(node.unpacked) do
            result = result .. padding .. "  " .. ix .. ". " .. tostring(guy) .. ",\n"
         end
         result = result .. padding .. "}"
      elseif node.dependencyGraph then
         result = result .. tostring(node.dependencyGraph)
      end
   elseif node.type == "constant" then
     result = result .. "<" .. node.constant .. ">"
   elseif node.type == "variable" then
      result = result .. "<" .. (node.name or "unnamed") .. ">"
   elseif node.type == "binding" then
      result = result .. "{" .. tostring(node.field) .. " -> "
      if node.constant then
         result = result .. tostring(node.constant.constant)
      elseif node.variable then
         result = result .. formatQueryNode(node.variable)
      end
      return result .. "}"
   elseif node.type == "object" then
      result = result .. "{"
      for _, binding in std.ipairs(node.bindings) do
         result = result .. formatQueryNode(binding) .. ", "
      end
      return result .. "}"
   elseif node.type == "mutate" then
      result = result .. "<" .. node.operator .. ">{"
      for _, binding in std.ipairs(node.bindings) do
         result = result .. formatQueryNode(binding) .. ", "
      end
      return result .. "}"
   elseif node.type == "union" or node.type == "choose" then
      result = result .. "{\n"
      for _, query in std.ipairs(node.queries) do
         result = result .. formatQueryNode(query, 1) .. ",\n"
      end
      return result .. "}"
   elseif node.type == "expression" then
      result = result .. " " .. node.operator .. "("
      for _, binding in std.ipairs(node.bindings) do
         result = result .. binding.field .. " = " .. formatQueryNode(binding.variable or binding.constant) .. ", "
      end
      result = string.sub(result, 1, -3) .. ")"
   end
   return result
end

local DefaultNodeMeta = {}
DefaultNodeMeta.__tostring = formatQueryNode

local function applyDefaultMeta(_, node)
   if type(node) == "table" and getmetatable(node) == nil then
      setmetatable(node, DefaultNodeMeta)
   end
end

-- Dependency Graph

DependencyGraph = {}

function DependencyGraph:new(obj)
   obj = obj or {}
   -- Essential state
   obj.unsorted = obj.unsorted or Set:new() -- set of nodes that need to be ordered
   obj.sorted = obj.sorted or {} -- append-only set of ordered nodes

   obj.unsatisfied = obj.unsatisfied or {} -- Number of terms required but not bound per node
   obj.dependents = obj.dependents or {} -- Sets of nodes depending on a term
   obj.strongDependents = obj.strongDependents or {} -- Sets of nodes strongly depending on (requiring stabilization of) a term
   obj.bound = obj.bound or Set:new() -- Set of terms bound by the currently ordered set of nodes
   obj.terms = obj.terms or Set:new() -- Set of all terms provided by any node in the graph
   obj.termGroups = obj.termGroups or Set:new() -- Set of sets of terms that are in some way joined

   setmetatable(obj, self)
   self.__index = self
   return obj
end

-- Get the variables this dgraph depends on for reification
function DependencyGraph:depends()
   local depends = Set:new()
   for term in std.pairs(self.dependents) do
      depends:add(term)
   end
   return depends / self.bound
end

function DependencyGraph:provided()
   return self.bound
end

function DependencyGraph:isSorted()
   return self.unsorted:length() == 0
end

function DependencyGraph:addObjectNode(node)
   local produces = Set:new()
   local depends = Set:new()
   for _, binding in std.ipairs(node.bindings or nothing) do
      if binding.variable then
         produces:add(binding.variable)
         depends:add(binding.variable)
      end
   end
   return self:add(node, depends, produces)
end

function DependencyGraph:addMutateNode(node)
   local produces = Set:new()
   local depends = Set:new()
   local weakDepends = Set:new()
   for _, binding in std.ipairs(node.bindings or nothing) do
      -- If the entity term isn't bound, the mutation produces it
      if binding.field == ENTITY_FIELD and not self.terms[binding.variable] then
         produces:add(binding.variable)
      end

      -- If the binding is bound on a variable that is produced in the query, it becomes a dependency of the query.
      if binding.variable then
         weakDepends:add(binding.variable)
      end
   end

   return self:add(node, depends, produces, weakDepends)
end

function DependencyGraph:addExpressionNode(node)
   -- TODO also need to consider projections and groupings as dependencies
   -- TODO handle productions other than just "return", this will require schemas
   local produces = Set:new()
   local depends = Set:new()
   local maybeDepends
   local strongDepend
   local anyDepends
   local hasConstant = false
   for _, binding in std.ipairs(node.bindings) do
     if binding.field == "return" and binding.variable then
       produces:add(binding.variable)
     elseif binding.variable then
       depends:add(binding.variable)
     elseif binding.constant then
       hasConstant = true
     end
   end
   if node.operator == "=" and produces:length() == 0 then
     anyDepends = not hasConstant and depends or Set:new()
     produces = depends:clone()
     depends = Set:new()
   end
   return self:add(node, depends, produces, maybeDepends, strongDepends, anyDepends)
end

function DependencyGraph:addSubqueryNode(node)
  local depends
  local provides = Set:new()

  local maybeDepends = Set:new()

  for _, var in std.pairs(node.outputs) do
    provides:add(var)
  end

  for _, body in std.ipairs(node.queries) do
    local subgraph = DependencyGraph:fromQueryGraph(body)
    maybeDepends:union(subgraph:depends(), true)
  end
  return self:add(node, depends, provides, maybeDepends)
end

function DependencyGraph:fromQueryGraph(query, terms, bound)
   local uniqueCounter = 0
   local dgraph = self
   if getmetatable(dgraph) ~= DependencyGraph then
      -- @FIXME: this naive injection strategy lets objects be scheduled willy-nilly after mutates...
      -- Not tracking this information is currently incidentally correct, but cannot be relied upon to order subqueries properly within their parents
      -- dgraph = self:new{terms = terms, bound = bound}
      dgraph = self:new()

   end
   dgraph.query = query
   query.dependencyGraph = dgraph

   util.walk(query, applyDefaultMeta)

   for _, node in std.ipairs(query.expressions or nothing) do
      dgraph:addExpressionNode(node)
   end
   for _, node in std.ipairs(query.objects or nothing) do
      dgraph:addObjectNode(node)
   end
   for _, node in std.ipairs(query.nots or nothing) do
      dgraph:addSubqueryNode(node)
   end
   for _, node in std.ipairs(query.unions or nothing) do
      dgraph:addSubqueryNode(node)
   end
   for _, node in std.ipairs(query.chooses or nothing) do
      dgraph:addSubqueryNode(node)
   end

   -- If the ENTITY variable of this mutate isn't already bound and it's a child of another mutate,
   -- then we need to uniquify it to prevent unification with other mutates.
   for _, node in std.ipairs(query.mutates or nothing) do
      if not dgraph.terms[node.variable] then
         if node.parent and node.parent.type == "mutate" then
            local old = node.variable
            node.variable = util.shallowCopy(node.variable)
            node.variable.name = "$$tmp" .. "-" .. node.variable.name .. "-" .. uniqueCounter
            uniqueCounter = uniqueCounter + 1
            for _, binding in std.pairs(node.bindings) do
               if binding.type == "binding" and binding.variable == old then
                  binding.variable = node.variable
               end
            end
            local parent = node.parent
            for _, binding in std.pairs(parent.bindings) do
               if binding.type == "binding" and binding.variable == old then
                  binding.variable = node.variable
               end
            end
         end
      end
   end
{'fraction_non_alphanumeric': 0.1025830258302583, 'fraction_numerical': 0.001021856372409878, 'mean_word_length': 2.705511148506521, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1407963', 'n_tokens_mistral': 5090, 'n_tokens_neox': 5003, 'n_words': 1649}